Give pgexecute a datatypes() method to get custom type names

Not used yet

2 files changed, 39 insertions, 1 deletions

diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 93841e12..14677e50 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -109,7 +109,6 @@ class PGExecute(object):
         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",
        pg_catalog.pg_encoding_to_char(d.encoding) as "Encoding",
@@ -119,6 +118,28 @@ class PGExecute(object):
     FROM pg_catalog.pg_database d
     ORDER BY 1;"""
 
+    datatypes_query = '''
+        SELECT n.nspname schema_name,
+               t.typname type_name
+        FROM   pg_catalog.pg_type t
+               INNER JOIN pg_catalog.pg_namespace n
+                  ON n.oid = t.typnamespace
+        WHERE ( t.typrelid = 0  -- non-composite types
+                OR (  -- composite type, but not a table
+                      SELECT c.relkind = 'c'
+                      FROM pg_catalog.pg_class c
+                      WHERE c.oid = t.typrelid
+                    )
+              )
+              AND NOT EXISTS( -- ignore array types
+                    SELECT  1
+                    FROM    pg_catalog.pg_type el
+                    WHERE   el.oid = t.typelem AND el.typarray = t.oid
+                  )
+              AND n.nspname <> 'pg_catalog'
+              AND n.nspname <> 'information_schema'
+        ORDER BY 1, 2;'''
+
     def __init__(self, database, user, password, host, port):
         self.dbname = database
         self.user = user
@@ -300,3 +321,12 @@ class PGExecute(object):
             cur.execute(self.functions_query)
             for row in cur:
                 yield row
+
+    def datatypes(self):
+        """Yields tuples of (schema_name, type_name)"""
+
+        with self.conn.cursor() as cur:
+            _logger.debug('Datatypes Query. sql: %r', self.datatypes_query)
+            cur.execute(self.datatypes_query)
+            for row in cur:
+                yield row
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index 4a82b5e4..02ddd578 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -71,6 +71,14 @@ def test_functions_query(executor):
     funcs = list(executor.functions())
     assert funcs == [('public', 'func1'), ('schema1', 'func2')]
 
+
+@dbtest
+def test_datatypes_query(executor):
+    run(executor, 'create type foo AS (a int, b text)')
+
+    types = list(executor.datatypes())
+    assert types == [('public', 'foo')]
+
 @dbtest
 def test_database_list(executor):
     databases = executor.databases()