Store function metadata (not used yet)

4 files changed, 52 insertions, 2 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/pgcompleter.py b/pgcli/pgcompleter.py
index 56131265..fe08a538 100644
--- a/pgcli/pgcompleter.py
+++ b/pgcli/pgcompleter.py
@@ -43,7 +43,7 @@ class PGCompleter(Completer):
 
         self.special_commands = []
         self.databases = []
-        self.dbmetadata = {'tables': {}}
+        self.dbmetadata = {'tables': {}, 'functions': {}}
         self.search_path = []
 
         self.all_completions = set(self.keywords + self.functions)
@@ -87,6 +87,11 @@ class PGCompleter(Completer):
         for schema in schemata:
             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)
 
     def extend_tables(self, table_data):
@@ -116,13 +121,27 @@ class PGCompleter(Completer):
 
         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 = {'tables': {}}
+        self.dbmetadata = {'tables': {}, 'functions': {}}
         self.all_completions = set(self.keywords)
 
     @staticmethod
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