diff options
| -rw-r--r-- | pgcli/pgcompleter.py | 71 | ||||
| -rw-r--r-- | tests/test_smart_completion_multiple_schemata.py | 13 | ||||
| -rw-r--r-- | tests/test_smart_completion_public_schema_only.py | 63 | ||||
| -rw-r--r-- | tests/test_sqlcompletion.py | 9 |
4 files changed, 129 insertions, 27 deletions
diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py index ec81f98e..337cf8f7 100644 --- a/pgcli/pgcompleter.py +++ b/pgcli/pgcompleter.py @@ -420,21 +420,34 @@ class PGCompleter(Completer): schema = self.escape_name(tbl.schema) relname = self.escape_name(tbl.name) - # We don't know if schema.relname is a table or view. Since - # tables and views cannot share the same name, we can check one - # at a time - try: - columns.extend(meta['tables'][schema][relname]) + if tbl.is_function: + # Return column names from a set-returning function + try: + # Get an array of FunctionMetadata objects + functions = meta['functions'][schema][relname] + except KeyError: + # No such function name + continue - # Table exists, so don't bother checking for a view - continue - except KeyError: - pass + for func in functions: + # func is a FunctionMetadata object + columns.extend(func.fieldnames()) + else: + # We don't know if schema.relname is a table or view. Since + # tables and views cannot share the same name, we can check + # one at a time + try: + columns.extend(meta['tables'][schema][relname]) - try: - columns.extend(meta['views'][schema][relname]) - except KeyError: - pass + # Table exists, so don't bother checking for a view + continue + except KeyError: + pass + + try: + columns.extend(meta['views'][schema][relname]) + except KeyError: + pass else: # Schema not specified, so traverse the search path looking for @@ -444,17 +457,27 @@ class PGCompleter(Completer): for schema in self.search_path: relname = self.escape_name(tbl.name) - try: - columns.extend(meta['tables'][schema][relname]) - break - except KeyError: - pass - - try: - columns.extend(meta['views'][schema][relname]) - break - except KeyError: - pass + if tbl.is_function: + try: + functions = meta['functions'][schema][relname] + except KeyError: + continue + + for func in functions: + # func is a FunctionMetadata object + columns.extend(func.fieldnames()) + else: + try: + columns.extend(meta['tables'][schema][relname]) + break + except KeyError: + pass + + try: + columns.extend(meta['views'][schema][relname]) + break + except KeyError: + pass return columns diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py index 6597f6d0..7eb81385 100644 --- a/tests/test_smart_completion_multiple_schemata.py +++ b/tests/test_smart_completion_multiple_schemata.py @@ -22,7 +22,9 @@ metadata = { ['func2', '', '', False, False, False]], 'custom': [ ['func3', '', '', False, False, False], - ['set_returning_func', '', '', False, False, True]], + ['set_returning_func', + 'OUT x INT', 'SETOF INT', + False, False, True]], }, 'datatypes': { 'public': ['typ1', 'typ2'], @@ -294,3 +296,12 @@ def test_schema_qualified_type_name(text, completer, complete_event): Completion(text='products', display_meta='table'), Completion(text='shipments', display_meta='table'), ]) + + +def test_suggest_columns_from_aliased_set_returning_function(completer, complete_event): + sql = 'select f. from custom.set_returning_func() f' + pos = len('select f.') + result = completer.get_completions(Document(text=sql, cursor_position=pos), + complete_event) + assert set(result) == set([ + Completion(text='x', start_position=0, display_meta='column')])
\ No newline at end of file diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py index cbf7d990..16e7a5ab 100644 --- a/tests/test_smart_completion_public_schema_only.py +++ b/tests/test_smart_completion_public_schema_only.py @@ -14,7 +14,8 @@ metadata = { 'functions': [ ['custom_func1', '', '', False, False, False], ['custom_func2', '', '', False, False, False], - ['set_returning_func', '', '', False, False, True]], + ['set_returning_func', '', 'TABLE (x INT, y INT)', + False, False, True]], 'datatypes': ['custom_type1', 'custom_type2'], } @@ -246,7 +247,7 @@ def test_suggested_multiple_column_names(completer, complete_event): Completion(text='last_name', start_position=0, display_meta='column'), Completion(text='custom_func1', start_position=0, display_meta='function'), Completion(text='custom_func2', start_position=0, display_meta='function'), - Completion(text='set_returning_func', start_position=0, display_meta='function')] + + Completion(text='set_returning_func', start_position=0, display_meta='function')] + list(map(lambda f: Completion(f, display_meta='function'), completer.functions)) + list(map(lambda x: Completion(x, display_meta='keyword'), completer.keywords)) ) @@ -416,3 +417,61 @@ def test_suggest_columns_from_escaped_table_alias(completer, complete_event): Completion(text='"insert"', start_position=0, display_meta='column'), Completion(text='"ABC"', start_position=0, display_meta='column'), ]) + + +def test_suggest_columns_from_set_returning_function(completer, complete_event): + sql = 'select from set_returning_func()' + pos = len('select ') + result = completer.get_completions(Document(text=sql, cursor_position=pos), + complete_event) + assert set(result) == set([ + Completion(text='x', start_position=0, display_meta='column'), + Completion(text='y', start_position=0, display_meta='column'), + Completion(text='custom_func1', start_position=0, display_meta='function'), + Completion(text='custom_func2', start_position=0, display_meta='function'), + Completion(text='set_returning_func', start_position=0, display_meta='function')] + + list(map(lambda f: Completion(f, display_meta='function'), completer.functions)) + + list(map(lambda x: Completion(x, display_meta='keyword'), completer.keywords))) + + +def test_suggest_columns_from_aliased_set_returning_function(completer, complete_event): + sql = 'select f. from set_returning_func() f' + pos = len('select f.') + result = completer.get_completions(Document(text=sql, cursor_position=pos), + complete_event) + assert set(result) == set([ + Completion(text='x', start_position=0, display_meta='column'), + Completion(text='y', start_position=0, display_meta='column')]) + + +def test_join_functions_using_suggests_common_columns(completer, complete_event): + text = '''SELECT * FROM set_returning_func() f1 + INNER JOIN set_returning_func() f2 USING (''' + pos = len(text) + result = set(completer.get_completions( + Document(text=text, cursor_position=pos), complete_event)) + assert set(result) == set([ + Completion(text='x', start_position=0, display_meta='column'), + Completion(text='y', start_position=0, display_meta='column')]) + + +def test_join_functions_using_suggests_common_columns(completer, complete_event): + text = '''SELECT * FROM set_returning_func() f1 + INNER JOIN set_returning_func() f2 USING (''' + pos = len(text) + result = set(completer.get_completions( + Document(text=text, cursor_position=pos), complete_event)) + assert set(result) == set([ + Completion(text='x', start_position=0, display_meta='column'), + Completion(text='y', start_position=0, display_meta='column')]) + + +def test_join_functions_on_suggests_columns(completer, complete_event): + text = '''SELECT * FROM set_returning_func() f1 + INNER JOIN set_returning_func() f2 ON f1.''' + pos = len(text) + result = set(completer.get_completions( + Document(text=text, cursor_position=pos), complete_event)) + assert set(result) == set([ + Completion(text='x', start_position=0, display_meta='column'), + Completion(text='y', start_position=0, display_meta='column')])
\ No newline at end of file diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py index 3842cd98..c2585336 100644 --- a/tests/test_sqlcompletion.py +++ b/tests/test_sqlcompletion.py @@ -612,3 +612,12 @@ def test_suggest_where_keyword(text): def test_named_query_completion(text, before, expected): suggestions = suggest_type(text, before) assert sorted_dicts(expected) == sorted_dicts(suggestions) + + +def test_select_suggests_fields_from_function(): + suggestions = suggest_type('SELECT FROM func()', 'SELECT ') + assert sorted_dicts(suggestions) == sorted_dicts([ + {'type': 'column', 'tables': [(None, 'func', None, True)]}, + {'type': 'function', 'schema': []}, + {'type': 'keyword'} + ]) |