Merge pull request #155 from darikg/functions

Autocomplete custom function names

8 files changed, 192 insertions, 70 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index 59af1336..4f73f05f 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -288,6 +288,7 @@ class PGCli(object):
         completer.extend_schemata(pgexecute.schemata())
         completer.extend_tables(pgexecute.tables())
         completer.extend_columns(pgexecute.columns())
+        completer.extend_functions(pgexecute.functions())
         completer.extend_database_names(pgexecute.databases())
 
     def get_completions(self, text, cursor_positition):
diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index 8c0d5fdb..c01277ed 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -95,10 +95,13 @@ def suggest_based_on_last_token(token, text_before_cursor, full_text):
         return [{'type': 'column', 'tables': extract_tables(full_text)}]
     elif token_v.lower() in ('select', 'where', 'having'):
         return [{'type': 'column', 'tables': extract_tables(full_text)},
-                {'type': 'function'}]
+                {'type': 'function', 'schema': []}]
     elif token_v.lower() in ('copy', 'from', 'update', 'into', 'describe',
             'join', 'table'):
         return [{'type': 'schema'}, {'type': 'table', 'schema': []}]
+    elif token_v.lower() == 'function':
+        # E.g. 'DROP FUNCTION <funcname>'
+        return [{'type': 'schema'}, {'type': 'function', 'schema': []}]
     elif token_v.lower() == 'on':
         tables = extract_tables(full_text)  # [(schema, table, alias), ...]
 
@@ -140,7 +143,8 @@ def suggest_based_on_last_token(token, text_before_cursor, full_text):
         suggestions.append({'type': 'column', 'tables': tables})
 
         # SCHEMA.<suggestion>
-        suggestions.append({'type': 'table', 'schema': identifier})
+        suggestions.extend([{'type': 'table', 'schema': identifier},
+                            {'type': 'function', 'schema': identifier}])
 
         return suggestions
 
diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py
index df36d2fd..35054cd1 100644
--- a/pgcli/pgcompleter.py
+++ b/pgcli/pgcompleter.py
@@ -43,7 +43,7 @@ class PGCompleter(Completer):
 
         self.special_commands = []
         self.databases = []
-        self.dbmetadata = {}
+        self.dbmetadata = {'tables': {}, 'functions': {}}
         self.search_path = []
 
         self.all_completions = set(self.keywords + self.functions)
@@ -81,10 +81,16 @@ class PGCompleter(Completer):
 
     def extend_schemata(self, schemata):
 
-        # data is a DataFrame with columns [schema]
+        # schemata is a list of schema names
         schemata = self.escaped_names(schemata)
+        metadata = self.dbmetadata['tables']
         for schema in schemata:
-            self.dbmetadata[schema] = {}
+            metadata[schema] = {}
+
+        # dbmetadata.values() are the 'tables' and 'functions' dicts
+        for metadata in self.dbmetadata.values():
+            for schema in schemata:
+                metadata[schema] = {}
 
         self.all_completions.update(schemata)
 
@@ -93,11 +99,12 @@ class PGCompleter(Completer):
         # table_data is a list of (schema_name, table_name) tuples
         table_data = [self.escaped_names(d) for d in table_data]
 
-        # dbmetadata['schema_name']['table_name'] should be a list of column
-        # names. Default to an asterisk
+        # dbmetadata['tables']['schema_name']['table_name'] should be a list of
+        # column names. Default to an asterisk
+        metadata = self.dbmetadata['tables']
         for schema, table in table_data:
             try:
-                self.dbmetadata[schema][table] = ['*']
+                metadata[schema][table] = ['*']
             except AttributeError:
                 _logger.error('Table %r listed in unrecognized schema %r',
                               table, schema)
@@ -108,19 +115,33 @@ class PGCompleter(Completer):
 
         # column_data is a list of (schema_name, table_name, column_name) tuples
         column_data = [self.escaped_names(d) for d in column_data]
-
+        metadata = self.dbmetadata['tables']
         for schema, table, column in column_data:
-            self.dbmetadata[schema][table].append(column)
+            metadata[schema][table].append(column)
 
         self.all_completions.update(t[2] for t in column_data)
 
+    def extend_functions(self, func_data):
+
+        # 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
+        metadata = self.dbmetadata['functions']
+        
+        for f in func_data:
+            schema, func = self.escaped_names(f)
+            metadata[schema][func] = None
+            self.all_completions.add(func)
+
     def set_search_path(self, search_path):
         self.search_path = self.escaped_names(search_path)
 
     def reset_completions(self):
         self.databases = []
         self.search_path = []
-        self.dbmetadata = {}
+        self.dbmetadata = {'tables': {}, 'functions': {}}
         self.all_completions = set(self.keywords)
 
     @staticmethod
@@ -156,30 +177,27 @@ class PGCompleter(Completer):
                 completions.extend(cols)
 
             elif suggestion['type'] == 'function':
-                funcs = self.find_matches(word_before_cursor, self.functions)
+                funcs = self.populate_schema_objects(
+                    suggestion['schema'], 'functions')
+
+                if not suggestion['schema']:
+                    # also suggest hardcoded functions
+                    funcs.extend(self.functions)
+
+                funcs = self.find_matches(word_before_cursor, funcs)
                 completions.extend(funcs)
 
             elif suggestion['type'] == 'schema':
-                schema_names = self.dbmetadata.keys()
+                schema_names = self.dbmetadata['tables'].keys()
                 schema_names = self.find_matches(word_before_cursor, schema_names)
                 completions.extend(schema_names)
 
             elif suggestion['type'] == 'table':
-
-                if suggestion['schema']:
-                    try:
-                        tables = self.dbmetadata[suggestion['schema']].keys()
-                    except KeyError:
-                        #schema doesn't exist
-                        tables = []
-                else:
-                    schemas = self.search_path
-                    meta = self.dbmetadata
-                    tables = [tbl for schema in schemas
-                                    for tbl in meta[schema].keys()]
-
+                tables = self.populate_schema_objects(
+                    suggestion['schema'], 'tables')
                 tables = self.find_matches(word_before_cursor, tables)
                 completions.extend(tables)
+
             elif suggestion['type'] == 'alias':
                 aliases = suggestion['aliases']
                 aliases = self.find_matches(word_before_cursor, aliases)
@@ -202,7 +220,7 @@ class PGCompleter(Completer):
         """
 
         columns = []
-        meta = self.dbmetadata
+        meta = self.dbmetadata['tables']
 
         for tbl in scoped_tbls:
             if tbl[0]:
@@ -226,5 +244,23 @@ class PGCompleter(Completer):
 
         return columns
 
+    def populate_schema_objects(self, schema, obj_type):
+        """Returns list of tables or functions for a (optional) schema"""
+
+        metadata = self.dbmetadata[obj_type]
+
+        if schema:
+            try:
+                objects = metadata[schema].keys()
+            except KeyError:
+                #schema doesn't exist
+                objects = []
+        else:
+            schemas = self.search_path
+            objects = [obj for schema in schemas
+                           for obj in metadata[schema].keys()]
+
+        return objects
+
 
 
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index e50b770d..b9fb62a7 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -63,6 +63,16 @@ class PGExecute(object):
                 AND att.attnum  > 0
         ORDER BY 1, 2, 3'''
 
+    functions_query = '''
+        SELECT 	DISTINCT  --multiple dispatch means possible duplicates
+                n.nspname schema_name,
+                p.proname func_name
+        FROM 	pg_catalog.pg_proc p
+                INNER JOIN pg_catalog.pg_namespace n
+                    ON n.oid = p.pronamespace
+        WHERE 	n.nspname NOT IN ('pg_catalog', 'information_schema')
+        ORDER BY 1, 2'''
+
 
     databases_query = """SELECT d.datname as "Name",
        pg_catalog.pg_get_userbyid(d.datdba) as "Owner",
@@ -186,3 +196,12 @@ class PGExecute(object):
             _logger.debug('Databases Query. sql: %r', self.databases_query)
             cur.execute(self.databases_query)
             return [x[0] for x in cur.fetchall()]
+
+    def functions(self):
+        """Yields tuples of (schema_name, function_name)"""
+
+        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
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index e9c97e36..6782e347 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -48,6 +48,17 @@ def test_schemata_table_and_columns_query(executor):
     assert executor.search_path() == ['public']
 
 @dbtest
+def test_functions_query(executor):
+    run(executor, '''create function func1() returns int
+                     language sql as $$select 1$$''')
+    run(executor, 'create schema schema1')
+    run(executor, '''create function schema1.func2() returns int
+                     language sql as $$select 2$$''')
+
+    funcs = list(executor.functions())
+    assert funcs == [('public', 'func1'), ('schema1', 'func2')]
+
+@dbtest
 def test_database_list(executor):
     databases = executor.databases()
     assert '_test_db' in databases
diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py
index d36f4c5a..66653f62 100644
--- a/tests/test_smart_completion_multiple_schemata.py
+++ b/tests/test_smart_completion_multiple_schemata.py
@@ -3,16 +3,20 @@ from prompt_toolkit.completion import Completion
 from prompt_toolkit.document import Document
 
 metadata = {
-            'public':   {
-                            'users': ['id', 'email', 'first_name', 'last_name'],
-                            'orders': ['id', 'ordered_date', 'status'],
-                            'select': ['id', 'insert', 'ABC']
-                        },
-            'custom':   {
-                            'users': ['id', 'phone_number'],
-                            'products': ['id', 'product_name', 'price'],
-                            'shipments': ['id', 'address', 'user_id']
-                        }
+            'tables': {
+                'public':   {
+                                'users': ['id', 'email', 'first_name', 'last_name'],
+                                'orders': ['id', 'ordered_date', 'status'],
+                                'select': ['id', 'insert', 'ABC']
+                            },
+                'custom':   {
+                                'users': ['id', 'phone_number'],
+                                'products': ['id', 'product_name', 'price'],
+                                'shipments': ['id', 'address', 'user_id']
+                            }},
+            'functions': {
+                'public':   ['func1', 'func2'],
+                'custom':   ['func3', 'func4']}
             }
 
 @pytest.fixture
@@ -23,16 +27,21 @@ def completer():
 
     schemata, tables, columns = [], [], []
 
-    for schema, tbls in metadata.items():
+    for schema, tbls in metadata['tables'].items():
         schemata.append(schema)
 
         for table, cols in tbls.items():
             tables.append((schema, table))
             columns.extend([(schema, table, col) for col in cols])
 
+    functions = [(schema, func)
+                    for schema, funcs in metadata['functions'].items()
+                    for func in funcs]
+
     comp.extend_schemata(schemata)
     comp.extend_tables(tables)
     comp.extend_columns(columns)
+    comp.extend_functions(functions)
     comp.set_search_path(['public'])
 
     return comp
@@ -70,7 +79,9 @@ def test_suggested_column_names_from_shadowed_visible_table(completer, complete_
         Completion(text='id', start_position=0),
         Completion(text='email', start_position=0),
         Completion(text='first_name', start_position=0),
-        Completion(text='last_name', start_position=0)] +
+        Completion(text='last_name', start_position=0),
+        Completion(text='func1', start_position=0),
+        Completion(text='func2', start_position=0)] +
         list(map(Completion, completer.functions)))
 
 def test_suggested_column_names_from_qualified_shadowed_table(completer, complete_event):
@@ -82,7 +93,9 @@ def test_suggested_column_names_from_qualified_shadowed_table(completer, complet
     assert set(result) == set([
         Completion(text='*', start_position=0),
         Completion(text='id', start_position=0),
-        Completion(text='phone_number', start_position=0)] +
+        Completion(text='phone_number', start_position=0),
+        Completion(text='func1', start_position=0),
+        Completion(text='func2', start_position=0)] +
         list(map(Completion, completer.functions)))
 
 def test_suggested_column_names_from_schema_qualifed_table(completer, complete_event):
@@ -100,7 +113,9 @@ def test_suggested_column_names_from_schema_qualifed_table(completer, complete_e
         Completion(text='*', start_position=0),
         Completion(text='id', start_position=0),
         Completion(text='product_name', start_position=0),
-        Completion(text='price', start_position=0)] +
+        Completion(text='price', start_position=0),
+        Completion(text='func1', start_position=0),
+        Completion(text='func2', start_position=0)] +
         list(map(Completion, completer.functions)))
 
 def test_suggested_column_names_in_function(completer, complete_event):
@@ -130,7 +145,9 @@ def test_suggested_table_names_with_schema_dot(completer, complete_event):
     assert set(result) == set([
         Completion(text='users', start_position=0),
         Completion(text='products', start_position=0),
-        Completion(text='shipments', start_position=0)])
+        Completion(text='shipments', start_position=0),
+        Completion(text='func3', start_position=0),
+        Completion(text='func4', start_position=0)])
 
 def test_suggested_column_names_with_qualified_alias(completer, complete_event):
     """
@@ -167,7 +184,9 @@ def test_suggested_multiple_column_names(completer, complete_event):
         Completion(text='*', start_position=0),
         Completion(text='id', start_position=0),
         Completion(text='product_name', start_position=0),
-        Completion(text='price', start_position=0)] +
+        Completion(text='price', start_position=0),
+        Completion(text='func1', start_position=0),
+        Completion(text='func2', start_position=0)] +
         list(map(Completion, completer.functions)))
 
 def test_suggested_multiple_column_names_with_alias(completer, complete_event):
@@ -223,5 +242,11 @@ def test_table_names_after_from(completer, complete_event):
         Completion(text='"select"', start_position=0),
         ])
 
-
-
+def test_schema_qualified_function_name(completer, complete_event):
+    text = 'SELECT custom.func'
+    postion = len(text)
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=postion), complete_event))
+    assert result == set([
+        Completion(text='func3', start_position=-len('func')),
+        Completion(text='func4', start_position=-len('func'))])
diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py
index 4c64c825..3f5d4fcc 100644
--- a/tests/test_smart_completion_public_schema_only.py
+++ b/tests/test_smart_completion_public_schema_only.py
@@ -3,9 +3,11 @@ from prompt_toolkit.completion import Completion
 from prompt_toolkit.document import Document
 
 metadata = {
-                'users': ['id', 'email', 'first_name', 'last_name'],
-                'orders': ['id', 'ordered_date', 'status'],
-                'select': ['id', 'insert', 'ABC']
+                'tables': {
+                    'users': ['id', 'email', 'first_name', 'last_name'],
+                    'orders': ['id', 'ordered_date', 'status'],
+                    'select': ['id', 'insert', 'ABC']},
+                'functions': ['custom_func1', 'custom_func2']
             }
 
 @pytest.fixture
@@ -17,13 +19,16 @@ def completer():
     schemata = ['public']
     tables, columns = [], []
 
-    for table, cols in metadata.items():
+    for table, cols in metadata['tables'].items():
         tables.append(('public', table))
         columns.extend([('public', table, col) for col in cols])
 
+    functions = [('public', func) for func in metadata['functions']]
+
     comp.extend_schemata(schemata)
     comp.extend_tables(tables)
     comp.extend_columns(columns)
+    comp.extend_functions(functions)
     comp.set_search_path(['public'])
 
     return comp
@@ -62,13 +67,22 @@ def test_schema_or_visible_table_completion(completer, complete_event):
                                Completion(text='orders', start_position=0)])
 
 
-def test_function_name_completion(completer, complete_event):
+def test_builtin_function_name_completion(completer, complete_event):
     text = 'SELECT MA'
     position = len('SELECT MA')
     result = completer.get_completions(
         Document(text=text, cursor_position=position), complete_event)
     assert set(result) == set([Completion(text='MAX', start_position=-2)])
 
+def test_user_function_name_completion(completer, complete_event):
+    text = 'SELECT cu'
+    position = len('SELECT cu')
+    result = completer.get_completions(
+        Document(text=text, cursor_position=position), complete_event)
+    assert set(result) == set([
+        Completion(text='custom_func1', start_position=-2),
+        Completion(text='custom_func2', start_position=-2)])
+
 def test_suggested_column_names_from_visible_table(completer, complete_event):
     """
     Suggest column and function names when selecting from table
@@ -86,7 +100,9 @@ def test_suggested_column_names_from_visible_table(completer, complete_event):
         Completion(text='id', start_position=0),
         Completion(text='email', start_position=0),
         Completion(text='first_name', start_position=0),
-        Completion(text='last_name', start_position=0)] +
+        Completion(text='last_name', start_position=0),
+        Completion(text='custom_func1', start_position=0),
+        Completion(text='custom_func2', start_position=0)] +
         list(map(Completion, completer.functions)))
 
 def test_suggested_column_names_in_function(completer, complete_event):
@@ -165,7 +181,9 @@ def test_suggested_multiple_column_names(completer, complete_event):
         Completion(text='id', start_position=0),
         Completion(text='email', start_position=0),
         Completion(text='first_name', start_position=0),
-        Completion(text='last_name', start_position=0)] +
+        Completion(text='last_name', start_position=0),
+        Completion(text='custom_func1', start_position=0),
+        Completion(text='custom_func2', start_position=0)] +
         list(map(Completion, completer.functions)))
 
 def test_suggested_multiple_column_names_with_alias(completer, complete_event):
@@ -271,5 +289,7 @@ def test_auto_escaped_col_names(completer, complete_event):
         Completion(text='*', start_position=0),
         Completion(text='id', start_position=0),
         Completion(text='"insert"', start_position=0),
-        Completion(text='"ABC"', start_position=0), ] +
+        Completion(text='"ABC"', start_position=0),
+        Completion(text='custom_func1', start_position=0),
+        Completion(text='custom_func2', start_position=0)] +
         list(map(Completion, completer.functions)))
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index 760e4dba..4c28d9d4 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -9,20 +9,20 @@ def test_select_suggests_cols_with_visible_table_scope():
     suggestions = suggest_type('SELECT  FROM tabl', 'SELECT ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'tabl', None)]},
-            {'type': 'function'}])
+            {'type': 'function', 'schema': []}])
 
 def test_select_suggests_cols_with_qualified_table_scope():
     suggestions = suggest_type('SELECT  FROM sch.tabl', 'SELECT ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [('sch', 'tabl', None)]},
-            {'type': 'function'}])
+            {'type': 'function', 'schema': []}])
 
 def test_where_suggests_columns_functions():
     suggestions = suggest_type('SELECT * FROM tabl WHERE ',
             'SELECT * FROM tabl WHERE ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'tabl', None)]},
-            {'type': 'function'}])
+            {'type': 'function', 'schema': []}])
 
 def test_lparen_suggests_cols():
     suggestion = suggest_type('SELECT MAX( FROM tbl', 'SELECT MAX(')
@@ -33,7 +33,7 @@ def test_select_suggests_cols_and_funcs():
     suggestions = suggest_type('SELECT ', 'SELECT ')
     assert sorted_dicts(suggestions) == sorted_dicts([
          {'type': 'column', 'tables': []},
-         {'type': 'function'}])
+         {'type': 'function', 'schema': []}])
 
 def test_from_suggests_tables_and_schemas():
     suggestions = suggest_type('SELECT * FROM ', 'SELECT * FROM ')
@@ -49,7 +49,7 @@ def test_col_comma_suggests_cols():
     suggestions = suggest_type('SELECT a, b, FROM tbl', 'SELECT a, b,')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'tbl', None)]},
-         {'type': 'function'}])
+        {'type': 'function', 'schema': []}])
 
 def test_table_comma_suggests_tables_and_schemas():
     suggestions = suggest_type('SELECT a, b FROM tbl1, ',
@@ -81,27 +81,30 @@ def test_partially_typed_col_name_suggests_col_names():
             'SELECT * FROM tabl WHERE col_n')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'tabl', None)]},
-        {'type': 'function'}])
+        {'type': 'function', 'schema': []}])
 
 def test_dot_suggests_cols_of_a_table_or_schema_qualified_table():
     suggestions = suggest_type('SELECT tabl. FROM tabl', 'SELECT tabl.')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'tabl', None)]},
-        {'type': 'table', 'schema': 'tabl'}])
+        {'type': 'table', 'schema': 'tabl'},
+        {'type': 'function', 'schema': 'tabl'}])
 
 def test_dot_suggests_cols_of_an_alias():
     suggestions = suggest_type('SELECT t1. FROM tabl1 t1, tabl2 t2',
             'SELECT t1.')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'table', 'schema': 't1'},
-        {'type': 'column', 'tables': [(None, 'tabl1', 't1')]}])
+        {'type': 'column', 'tables': [(None, 'tabl1', 't1')]},
+        {'type': 'function', 'schema': 't1'}])
 
 def test_dot_col_comma_suggests_cols_or_schema_qualified_table():
     suggestions = suggest_type('SELECT t1.a, t2. FROM tabl1 t1, tabl2 t2',
             'SELECT t1.a, t2.')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'tabl2', 't2')]},
-        {'type': 'table', 'schema': 't2'}])
+        {'type': 'table', 'schema': 't2'},
+        {'type': 'function', 'schema': 't2'}])
 
 def test_sub_select_suggests_keyword():
     suggestion = suggest_type('SELECT * FROM (', 'SELECT * FROM (')
@@ -122,7 +125,7 @@ def test_sub_select_col_name_completion():
             'SELECT * FROM (SELECT ')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'abc', None)]},
-        {'type': 'function'}])
+        {'type': 'function', 'schema': []}])
 
 @pytest.mark.xfail
 def test_sub_select_multiple_col_name_completion():
@@ -130,14 +133,15 @@ def test_sub_select_multiple_col_name_completion():
             'SELECT * FROM (SELECT a, ')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'abc', None)]},
-        {'type': 'function'}])
+        {'type': 'function', 'schema': []}])
 
 def test_sub_select_dot_col_name_completion():
     suggestions = suggest_type('SELECT * FROM (SELECT t. FROM tabl t',
             'SELECT * FROM (SELECT t.')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'tabl', 't')]},
-        {'type': 'table', 'schema': 't'}])
+        {'type': 'table', 'schema': 't'},
+        {'type': 'function', 'schema': 't'}])
 
 def test_join_suggests_tables_and_schemas():
     suggestion = suggest_type('SELECT * FROM abc a JOIN ',
@@ -151,14 +155,16 @@ def test_join_alias_dot_suggests_cols1():
             'SELECT * FROM abc a JOIN def d ON a.')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'abc', 'a')]},
-        {'type': 'table', 'schema': 'a'}])
+        {'type': 'table', 'schema': 'a'},
+        {'type': 'function', 'schema': 'a'}])
 
 def test_join_alias_dot_suggests_cols2():
     suggestion = suggest_type('SELECT * FROM abc a JOIN def d ON a.',
             'SELECT * FROM abc a JOIN def d ON a.id = d.')
     assert sorted_dicts(suggestion) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'def', 'd')]},
-        {'type': 'table', 'schema': 'd'}])
+        {'type': 'table', 'schema': 'd'},
+        {'type': 'function', 'schema': 'd'}])
 
 def test_on_suggests_aliases():
     suggestions = suggest_type(
@@ -194,7 +200,7 @@ def test_2_statements_2nd_current():
                                'select * from a; select ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'b', None)]},
-            {'type': 'function'}])
+            {'type': 'function', 'schema': []}])
 
     # Should work even if first statement is invalid
     suggestions = suggest_type('select * from; select * from ',
@@ -212,7 +218,7 @@ def test_2_statements_1st_current():
                                'select ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'a', None)]},
-            {'type': 'function'}])
+            {'type': 'function', 'schema': []}])
 
 def test_3_statements_2nd_current():
     suggestions = suggest_type('select * from a; select * from ; select * from c',
@@ -224,5 +230,5 @@ def test_3_statements_2nd_current():
                                'select * from a; select ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'b', None)]},
-            {'type': 'function'}])
+            {'type': 'function', 'schema': []}])