mirror of
https://github.com/dbcli/pgcli.git
synced 2024-10-06 10:17:15 +03:00
Make suggestions based on local tables
This commit is contained in:
parent
d113d4e38f
commit
2407de0e4c
@ -366,7 +366,8 @@ class PGCompleter(Completer):
|
||||
def get_column_matches(self, suggestion, word_before_cursor):
|
||||
tables = suggestion.table_refs
|
||||
_logger.debug("Completion column scope: %r", tables)
|
||||
scoped_cols = self.populate_scoped_cols(tables)
|
||||
scoped_cols = self.populate_scoped_cols(tables, suggestion.local_tables)
|
||||
|
||||
colit = scoped_cols.items
|
||||
flat_cols = list(chain(*((c.name for c in cols)
|
||||
for t, cols in colit())))
|
||||
@ -567,6 +568,7 @@ class PGCompleter(Completer):
|
||||
|
||||
def get_table_matches(self, suggestion, word_before_cursor, alias=False):
|
||||
tables = self.populate_schema_objects(suggestion.schema, 'tables')
|
||||
tables.extend(tbl.name for tbl in suggestion.local_tables)
|
||||
|
||||
# Unless we're sure the user really wants them, don't suggest the
|
||||
# pg_catalog tables that are implicitly on the search path
|
||||
@ -654,9 +656,10 @@ class PGCompleter(Completer):
|
||||
Path: get_path_matches,
|
||||
}
|
||||
|
||||
def populate_scoped_cols(self, scoped_tbls):
|
||||
def populate_scoped_cols(self, scoped_tbls, local_tbls=()):
|
||||
""" Find all columns in a set of scoped_tables
|
||||
:param scoped_tbls: list of TableReference namedtuples
|
||||
:param local_tbls: tuple(TableMetadata)
|
||||
:return: {TableReference:{colname:ColumnMetaData}}
|
||||
"""
|
||||
|
||||
@ -689,6 +692,10 @@ class PGCompleter(Completer):
|
||||
cols = cols.values()
|
||||
addcols(schema, relname, tbl.alias, reltype, cols)
|
||||
break
|
||||
|
||||
# Local tables should shadow database tables
|
||||
for tbl in local_tbls:
|
||||
columns[tbl.name] = tbl.columns
|
||||
|
||||
return columns
|
||||
|
||||
|
@ -3,6 +3,7 @@ import pytest
|
||||
from metadata import (MetaData, alias, name_join, fk_join, join, keyword,
|
||||
schema, table, view, function, column, wildcard_expansion)
|
||||
from prompt_toolkit.document import Document
|
||||
from prompt_toolkit.completion import Completion
|
||||
|
||||
metadata = {
|
||||
'tables': {
|
||||
@ -867,3 +868,49 @@ def test_insert(completer, complete_event, text):
|
||||
result = completer.get_completions(Document(text=text, cursor_position=pos),
|
||||
complete_event)
|
||||
assert set(result) == set(testdata.columns('users'))
|
||||
|
||||
|
||||
def test_suggest_cte_names(completer, complete_event):
|
||||
text = '''
|
||||
WITH cte1 AS (SELECT a, b, c FROM foo),
|
||||
cte2 AS (SELECT d, e, f FROM bar)
|
||||
SELECT * FROM
|
||||
'''
|
||||
pos = len(text)
|
||||
result = completer.get_completions(
|
||||
Document(text=text, cursor_position=pos),
|
||||
complete_event)
|
||||
expected = set([
|
||||
Completion('cte1', 0, display_meta='table'),
|
||||
Completion('cte2', 0, display_meta='table'),
|
||||
])
|
||||
assert expected <= set(result)
|
||||
|
||||
|
||||
def test_suggest_columns_from_cte(completer, complete_event):
|
||||
text = 'WITH cte AS (SELECT foo, bar FROM baz) SELECT FROM cte'
|
||||
pos = len('WITH cte AS (SELECT foo, bar FROM baz) SELECT ')
|
||||
result = completer.get_completions(Document(text=text, cursor_position=pos),
|
||||
complete_event)
|
||||
expected = ([Completion('foo', 0, display_meta='column'),
|
||||
Completion('bar', 0, display_meta='column'),
|
||||
] +
|
||||
testdata.functions() +
|
||||
testdata.builtin_functions() +
|
||||
testdata.keywords()
|
||||
)
|
||||
|
||||
assert set(expected) == set(result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('text', [
|
||||
'WITH cte AS (SELECT foo FROM bar) SELECT * FROM cte WHERE cte.',
|
||||
'WITH cte AS (SELECT foo FROM bar) SELECT * FROM cte c WHERE c.',
|
||||
])
|
||||
def test_cte_qualified_columns(completer, complete_event, text):
|
||||
pos = len(text)
|
||||
result = completer.get_completions(
|
||||
Document(text=text, cursor_position=pos),
|
||||
complete_event)
|
||||
expected = [Completion('foo', 0, display_meta='column')]
|
||||
assert set(expected) == set(result)
|
||||
|
Loading…
Reference in New Issue
Block a user