diff options
| -rw-r--r-- | pgcli/pgcompleter.py | 20 | ||||
| -rw-r--r-- | pgcli/pgexecute.py | 21 | ||||
| -rw-r--r-- | tests/test_pgexecute.py | 19 | ||||
| -rw-r--r-- | tests/test_smart_completion_multiple_schemata.py | 3 | ||||
| -rw-r--r-- | tests/test_smart_completion_public_schema_only.py | 4 |
5 files changed, 52 insertions, 15 deletions
diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py index 6281b08b..8b0b94c3 100644 --- a/pgcli/pgcompleter.py +++ b/pgcli/pgcompleter.py @@ -139,17 +139,23 @@ class PGCompleter(Completer): self.all_completions.add(column) def extend_functions(self, func_data): + + # func_data is a list of function metadata namedtuples + # with fields schema_name, func_name, arg_list, result, + # is_aggregate, is_window, is_set_returning - # func_data is an iterator of (schema_name, function_name) - - # dbmetadata['functions']['schema_name']['function_name'] should return - # function metadata -- right now we're not storing any further metadata - # so just default to None as a placeholder + # dbmetadata['schema_name']['functions']['function_name'] should return + # the function metadata namedtuple for the corresponding function metadata = self.dbmetadata['functions'] for f in func_data: - schema, func = self.escaped_names(f) - metadata[schema][func] = None + schema, func = self.escaped_names([f.schema_name, f.func_name]) + + if func in metadata[schema]: + metadata[schema][func].append(f) + else: + metadata[schema][func] = [f] + self.all_completions.add(func) def extend_datatypes(self, type_data): diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 0c3c356f..9a03f0eb 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -4,10 +4,12 @@ import psycopg2 import psycopg2.extras import psycopg2.extensions as ext import sqlparse +from collections import namedtuple from .packages import pgspecial as special from .encodingutils import unicode2utf8, PY2, utf8tounicode import click + _logger = logging.getLogger(__name__) # Cast all database input to unicode automatically. @@ -24,6 +26,10 @@ ext.register_type(ext.new_type((17,), 'BYTEA_TEXT', psycopg2.STRING)) # See http://initd.org/psycopg/articles/2014/07/20/cancelling-postgresql-statements-python/ ext.set_wait_callback(psycopg2.extras.wait_select) +FunctionMetadata = namedtuple('FunctionMetadata', + ['schema_name', 'func_name', 'arg_list', 'result', + 'is_aggregate', 'is_window', 'is_set_returning']) + def register_json_typecasters(conn, loads_fn): """Set the function for converting JSON data for a connection. @@ -102,14 +108,19 @@ class PGExecute(object): ORDER BY 1, 2, 3''' functions_query = ''' - SELECT DISTINCT --multiple dispatch means possible duplicates - n.nspname schema_name, - p.proname func_name + SELECT n.nspname schema_name, + p.proname func_name, + pg_catalog.pg_get_function_arguments(p.oid) arg_list, + pg_catalog.pg_get_function_result(p.oid) result, + p.proisagg is_aggregate, + p.proiswindow is_window, + p.proretset is_set_returning FROM pg_catalog.pg_proc p INNER JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace ORDER BY 1, 2''' + databases_query = """SELECT d.datname as "Name", pg_catalog.pg_get_userbyid(d.datdba) as "Owner", pg_catalog.pg_encoding_to_char(d.encoding) as "Encoding", @@ -346,13 +357,13 @@ class PGExecute(object): return [x[0] for x in cur.fetchall()] def functions(self): - """Yields tuples of (schema_name, function_name)""" + """Yields FunctionMetadata named tuples""" with self.conn.cursor() as cur: _logger.debug('Functions Query. sql: %r', self.functions_query) cur.execute(self.functions_query) for row in cur: - yield row + yield FunctionMetadata(*row) def datatypes(self): """Yields tuples of (schema_name, type_name)""" diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py index 3f6d4e86..dc289220 100644 --- a/tests/test_pgexecute.py +++ b/tests/test_pgexecute.py @@ -4,6 +4,7 @@ import pytest from pgcli.packages.pgspecial import PGSpecial from textwrap import dedent from utils import run, dbtest, requires_json, requires_jsonb +from pgcli.pgexecute import FunctionMetadata @dbtest def test_conn(executor): @@ -68,8 +69,24 @@ def test_functions_query(executor): run(executor, '''create function schema1.func2() returns int language sql as $$select 2$$''') + run(executor, '''create function func3() + returns table(x int, y int) language sql + as $$select 1, 2 from generate_series(1,5)$$;''') + + run(executor, '''create function func4(x int) returns setof int language sql + as $$select generate_series(1,5)$$;''') + funcs = set(executor.functions()) - assert funcs >= set([('public', 'func1'), ('schema1', 'func2')]) + assert funcs >= set([ + FunctionMetadata('public', 'func1', '', + 'integer', False, False, False), + FunctionMetadata('public', 'func3', '', + 'TABLE(x integer, y integer)', False, False, True), + FunctionMetadata('public', 'func4', 'x integer', + 'SETOF integer', False, False, True), + FunctionMetadata('schema1', 'func2', '', + 'integer', False, False, False), + ]) @dbtest diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py index fd8854d6..3a80a42f 100644 --- a/tests/test_smart_completion_multiple_schemata.py +++ b/tests/test_smart_completion_multiple_schemata.py @@ -2,6 +2,7 @@ from __future__ import unicode_literals import pytest from prompt_toolkit.completion import Completion from prompt_toolkit.document import Document +from pgcli.pgexecute import FunctionMetadata metadata = { 'tables': { @@ -40,7 +41,7 @@ def completer(): tables.append((schema, table)) columns.extend([(schema, table, col) for col in cols]) - functions = [(schema, func) + functions = [FunctionMetadata(schema, func, '', '', False, False, False) for schema, funcs in metadata['functions'].items() for func in funcs] diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py index 2f75069d..e1e2ac17 100644 --- a/tests/test_smart_completion_public_schema_only.py +++ b/tests/test_smart_completion_public_schema_only.py @@ -2,6 +2,7 @@ from __future__ import unicode_literals import pytest from prompt_toolkit.completion import Completion from prompt_toolkit.document import Document +from pgcli.pgexecute import FunctionMetadata metadata = { 'tables': { @@ -42,7 +43,8 @@ def completer(): comp.extend_columns(columns, kind='views') # functions - functions = [('public', func) for func in metadata['functions']] + functions = [FunctionMetadata('public', func, '', '', False, False, False) + for func in metadata['functions']] comp.extend_functions(functions) # types |