mirrors
/
pgcli
mirror of https://github.com/dbcli/pgcli
1
0
Fork 0
Code

Merge pull request #361 from dbcli/darikg/suggest-functions-as-tables 

Suggest functions as tables
This commit is contained in:
Amjith Ramanujam 2015-09-24 22:35:40 -07:00
parent 6166291c3e a1785a8aaa
commit 764f69a47b
7 changed files with 171 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
pgcli/packages/sqlcompletion.py
View File

@ -261,6 +261,7 @@ def suggest_based_on_last_token(token, text_before_cursor, full_text, identifier
{'type': 'keyword'}]
elif (token_v.endswith('join') and token.is_keyword) or (token_v in
('copy', 'from', 'update', 'into', 'describe', 'truncate')):
schema = (identifier and identifier.get_parent_name()) or []
# Suggest tables from either the currently-selected schema or the
@ -274,6 +275,11 @@ def suggest_based_on_last_token(token, text_before_cursor, full_text, identifier
# Only tables can be TRUNCATED, otherwise suggest views
if token_v != 'truncate':
suggest.append({'type': 'view', 'schema': schema})
# Suggest set-returning functions in the FROM clause
if token_v == 'from' or (token_v.endswith('join') and token.is_keyword):
suggest.append({'type': 'function', 'schema': schema,
'filter': 'is_set_returning'})
return suggest

74
pgcli/pgcompleter.py
View File

@ -2,6 +2,7 @@ from __future__ import print_function, unicode_literals
import logging
import re
import itertools
import operator
from prompt_toolkit.completion import Completer, Completion
from .packages.sqlcompletion import suggest_type
from .packages.parseutils import last_word
@ -42,6 +43,7 @@ class PGCompleter(Completer):
functions = ['AVG', 'COUNT', 'FIRST', 'FORMAT', 'LAST', 'LCASE', 'LEN',
'MAX', 'MIN', 'MID', 'NOW', 'ROUND', 'SUM', 'TOP', 'UCASE']
datatypes = ['BIGINT', 'BOOLEAN', 'CHAR', 'DATE', 'DOUBLE PRECISION', 'INT',
'INTEGER', 'NUMERIC', 'REAL', 'TEXT', 'VARCHAR']
@ -139,17 +141,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):
@ -271,14 +279,23 @@ class PGCompleter(Completer):
completions.extend(cols)
elif suggestion['type'] == 'function':
# suggest user-defined functions using substring matching
funcs = self.populate_schema_objects(
suggestion['schema'], 'functions')
user_funcs = self.find_matches(word_before_cursor, funcs,
meta='function')
completions.extend(user_funcs)
if suggestion.get('filter') == 'is_set_returning':
# Only suggest set-returning functions
filt = operator.attrgetter('is_set_returning')
funcs = self.populate_functions(suggestion['schema'], filt)
else:
funcs = self.populate_schema_objects(
suggestion['schema'], 'functions')
if not suggestion['schema']:
# Function overloading means we way have multiple functions
# of the same name at this point, so keep unique names only
funcs = set(funcs)
funcs = self.find_matches(word_before_cursor, funcs,
meta='function')
completions.extend(funcs)
if not suggestion['schema'] and 'filter' not in suggestion:
# also suggest hardcoded functions using startswith
# matching
predefined_funcs = self.find_matches(word_before_cursor,
@ -454,3 +471,32 @@ class PGCompleter(Completer):
for obj in metadata[schema].keys()]
return objects
def populate_functions(self, schema, filter_func):
"""Returns a list of function names
filter_func is a function that accepts a FunctionMetadata namedtuple
and returns a boolean indicating whether that function should be
kept or discarded
"""
metadata = self.dbmetadata['functions']
# Because of multiple dispatch, we can have multiple functions
# with the same name, which is why `for meta in metas` is necessary
# in the comprehensions below
if schema:
try:
return [func for (func, metas) in metadata[schema].items()
for meta in metas
if filter_func(meta)]
except KeyError:
return []
else:
return [func for schema in self.search_path
for (func, metas) in metadata[schema].items()
for meta in metas
if filter_func(meta)]

21
pgcli/pgexecute.py
View File

@ -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)"""

19
tests/test_pgexecute.py
View File

@ -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

20
tests/test_smart_completion_multiple_schemata.py
View File

@ -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': {
@ -16,8 +17,12 @@ metadata = {
'shipments': ['id', 'address', 'user_id']
}},
'functions': {
'public': ['func1', 'func2'],
'custom': ['func3', 'func4'],
'public': [
['func1', '', '', False, False, False],
['func2', '', '', False, False, False]],
'custom': [
['func3', '', '', False, False, False],
['set_returning_func', '', '', False, False, True]],
},
'datatypes': {
'public': ['typ1', 'typ2'],
@ -40,9 +45,9 @@ def completer():
tables.append((schema, table))
columns.extend([(schema, table, col) for col in cols])
functions = [(schema, func)
functions = [FunctionMetadata(schema, *func_meta)
for schema, funcs in metadata['functions'].items()
for func in funcs]
for func_meta in funcs]
datatypes = [(schema, datatype)
for schema, datatypes in metadata['datatypes'].items()
@ -164,8 +169,9 @@ def test_suggested_table_names_with_schema_dot(completer, complete_event):
assert set(result) == set([
Completion(text='users', start_position=0, display_meta='table'),
Completion(text='products', start_position=0, display_meta='table'),
Completion(text='shipments', start_position=0, display_meta='table')])
Completion(text='shipments', start_position=0, display_meta='table'),
Completion(text='set_returning_func', start_position=0, display_meta='function'),
])
def test_suggested_column_names_with_qualified_alias(completer, complete_event):
"""
Suggest column names on table alias and dot
@ -268,7 +274,7 @@ def test_schema_qualified_function_name(completer, complete_event):
Document(text=text, cursor_position=postion), complete_event))
assert result == set([
Completion(text='func3', start_position=-len('func'), display_meta='function'),
Completion(text='func4', start_position=-len('func'), display_meta='function')])
Completion(text='set_returning_func', start_position=-len('func'), display_meta='function')])
@pytest.mark.parametrize('text', [

22
tests/test_smart_completion_public_schema_only.py
View File

@ -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': {
@ -10,7 +11,10 @@ metadata = {
'select': ['id', 'insert', 'ABC']},
'views': {
'user_emails': ['id', 'email']},
'functions': ['custom_func1', 'custom_func2'],
'functions': [
['custom_func1', '', '', False, False, False],
['custom_func2', '', '', False, False, False],
['set_returning_func', '', '', False, False, True]],
'datatypes': ['custom_type1', 'custom_type2'],
}
@ -42,7 +46,8 @@ def completer():
comp.extend_columns(columns, kind='views')
# functions
functions = [('public', func) for func in metadata['functions']]
functions = [FunctionMetadata('public', *func_meta)
for func_meta in metadata['functions']]
comp.extend_functions(functions)
# types
@ -88,7 +93,8 @@ def test_schema_or_visible_table_completion(completer, complete_event):
Completion(text='users', start_position=0, display_meta='table'),
Completion(text='"select"', start_position=0, display_meta='table'),
Completion(text='orders', start_position=0, display_meta='table'),
Completion(text='user_emails', start_position=0, display_meta='view')])
Completion(text='user_emails', start_position=0, display_meta='view'),
Completion(text='set_returning_func', start_position=0, display_meta='function')])
def test_builtin_function_name_completion(completer, complete_event):
@ -154,7 +160,8 @@ def test_suggested_column_names_from_visible_table(completer, complete_event):
Completion(text='first_name', start_position=0, display_meta='column'),
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='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))
)
@ -238,7 +245,8 @@ def test_suggested_multiple_column_names(completer, complete_event):
Completion(text='first_name', start_position=0, display_meta='column'),
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='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))
)
@ -355,6 +363,7 @@ def test_table_names_after_from(completer, complete_event):
Completion(text='orders', start_position=0, display_meta='table'),
Completion(text='"select"', start_position=0, display_meta='table'),
Completion(text='user_emails', start_position=0, display_meta='view'),
Completion(text='set_returning_func', start_position=0, display_meta='function')
])
def test_auto_escaped_col_names(completer, complete_event):
@ -369,7 +378,8 @@ def test_auto_escaped_col_names(completer, complete_event):
Completion(text='"insert"', start_position=0, display_meta='column'),
Completion(text='"ABC"', 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='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))
)

51
tests/test_sqlcompletion.py
View File

@ -79,39 +79,64 @@ def test_select_suggests_cols_and_funcs():
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'column', 'tables': []},
{'type': 'function', 'schema': []},
{'type': 'keyword'}
{'type': 'keyword'},
])
@pytest.mark.parametrize('expression', [
'SELECT * FROM ',
'INSERT INTO ',
'COPY ',
'UPDATE ',
'DESCRIBE ',
'SELECT * FROM foo JOIN ',
])
def test_expression_suggests_tables_views_and_schemas(expression):
def test_suggests_tables_views_and_schemas(expression):
suggestions = suggest_type(expression, expression)
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'schema'}])
{'type': 'schema'},
])
@pytest.mark.parametrize('expression', [
'SELECT * FROM ',
'SELECT * FROM foo JOIN ',
])
def test_suggest_tables_views_schemas_and_set_returning_functions(expression):
suggestions = suggest_type(expression, expression)
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'function', 'schema': [], 'filter': 'is_set_returning'},
{'type': 'schema'},
])
@pytest.mark.parametrize('expression', [
'SELECT * FROM sch.',
'INSERT INTO sch.',
'COPY sch.',
'UPDATE sch.',
'DESCRIBE sch.',
'SELECT * FROM foo JOIN sch.',
])
def test_expression_suggests_qualified_tables_views_and_schemas(expression):
def test_suggest_qualified_tables_and_views(expression):
suggestions = suggest_type(expression, expression)
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': 'sch'},
{'type': 'view', 'schema': 'sch'}])
{'type': 'view', 'schema': 'sch'},
])
@pytest.mark.parametrize('expression', [
'SELECT * FROM sch.',
'SELECT * FROM foo JOIN sch.',
])
def test_suggest_qualified_tables_views_and_set_returning_functions(expression):
suggestions = suggest_type(expression, expression)
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': 'sch'},
{'type': 'view', 'schema': 'sch'},
{'type': 'function', 'schema': 'sch', 'filter': 'is_set_returning'},
])
def test_truncate_suggests_tables_and_schemas():
@ -147,6 +172,7 @@ def test_table_comma_suggests_tables_and_schemas():
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'function', 'schema': [], 'filter': 'is_set_returning'},
{'type': 'schema'}])
@ -272,6 +298,7 @@ def test_sub_select_table_name_completion(expression):
assert sorted_dicts(suggestion) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'function', 'schema': [], 'filter': 'is_set_returning'},
{'type': 'schema'}])
@ -314,6 +341,7 @@ def test_join_suggests_tables_and_schemas(tbl_alias, join_type):
assert sorted_dicts(suggestion) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'function', 'schema': [], 'filter': 'is_set_returning'},
{'type': 'schema'}])
@ -323,6 +351,7 @@ def test_left_join_with_comma():
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'function', 'schema': [], 'filter': 'is_set_returning'},
{'type': 'schema'}])
@ -389,6 +418,7 @@ def test_2_statements_2nd_current():
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'function', 'schema': [], 'filter': 'is_set_returning'},
{'type': 'schema'}])
suggestions = suggest_type('select * from a; select from b',
@ -405,6 +435,7 @@ def test_2_statements_2nd_current():
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'function', 'schema': [], 'filter': 'is_set_returning'},
{'type': 'schema'}])
@ -414,6 +445,7 @@ def test_2_statements_1st_current():
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'function', 'schema': [], 'filter': 'is_set_returning'},
{'type': 'schema'}])
suggestions = suggest_type('select from a; select * from b',
@ -431,6 +463,7 @@ def test_3_statements_2nd_current():
assert sorted_dicts(suggestions) == sorted_dicts([
{'type': 'table', 'schema': []},
{'type': 'view', 'schema': []},
{'type': 'function', 'schema': [], 'filter': 'is_set_returning'},
{'type': 'schema'}])
suggestions = suggest_type('select * from a; select from b; select * from c',