Merge branch 'master' of https://github.com/rclmenezes/sqlalchemy

1 files changed, 70 insertions, 21 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index b3506f5d2..19d2c7ca4 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -401,6 +401,23 @@ The value passed to the keyword argument will be simply passed through to the
 underlying CREATE INDEX command, so it *must* be a valid index type for your
 version of PostgreSQL.
 
+Special Reflection Options
+--------------------------
+
+The :class:`.Inspector` used for the Postgresql backend is an instance
+of :class:`.PGInspector`, which offers additional methods::
+
+    from sqlalchemy import create_engine, inspect
+
+    engine = create_engine("postgresql+psycopg2://localhost/test")
+    insp = inspect(engine)  # will be a PGInspector
+
+    print(insp.get_enums())
+
+.. autoclass:: PGInspector
+    :members:
+
+
 """
 from collections import defaultdict
 import re
@@ -1570,11 +1587,32 @@ class PGInspector(reflection.Inspector):
         reflection.Inspector.__init__(self, conn)
 
     def get_table_oid(self, table_name, schema=None):
-        """Return the oid from `table_name` and `schema`."""
+        """Return the OID for the given table name."""
 
         return self.dialect.get_table_oid(self.bind, table_name, schema,
                                           info_cache=self.info_cache)
 
+    def get_enums(self, schema=None):
+        """Return a list of ENUM objects.
+
+        Each member is a dictionary containing these fields:
+
+            * name - name of the enum
+            * schema - the schema name for the enum.
+            * visible - boolean, whether or not this enum is visible
+              in the default search path.
+            * labels - a list of string labels that apply to the enum.
+
+        :param schema: schema name.  If None, the default schema
+         (typically 'public') is used.  May also be set to '*' to
+         indicate load enums for all schemas.
+
+        .. versionadded:: 1.0.0
+
+        """
+        schema = schema or self.default_schema_name
+        return self.dialect._load_enums(self.bind, schema)
+
 
 class CreateEnumType(schema._CreateDropBase):
     __visit_name__ = "create_enum_type"
@@ -2056,7 +2094,12 @@ class PGDialect(default.DefaultDialect):
         c = connection.execute(s, table_oid=table_oid)
         rows = c.fetchall()
         domains = self._load_domains(connection)
-        enums = self._load_enums(connection)
+        enums = dict(
+            (
+                "%s.%s" % (rec['schema'], rec['name'])
+                if not rec['visible'] else rec['name'], rec) for rec in
+            self._load_enums(connection, schema='*')
+        )
 
         # format columns
         columns = []
@@ -2130,10 +2173,9 @@ class PGDialect(default.DefaultDialect):
             elif attype in enums:
                 enum = enums[attype]
                 coltype = ENUM
-                if "." in attype:
-                    kwargs['schema'], kwargs['name'] = attype.split('.')
-                else:
-                    kwargs['name'] = attype
+                kwargs['name'] = enum['name']
+                if not enum['visible']:
+                    kwargs['schema'] = enum['schema']
                 args = tuple(enum['labels'])
                 break
             elif attype in domains:
@@ -2444,7 +2486,8 @@ class PGDialect(default.DefaultDialect):
             for name, uc in uniques.items()
         ]
 
-    def _load_enums(self, connection):
+    def _load_enums(self, connection, schema=None):
+        schema = schema or self.default_schema_name
         if not self.supports_native_enum:
             return {}
 
@@ -2460,31 +2503,37 @@ class PGDialect(default.DefaultDialect):
                  LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
                  LEFT JOIN pg_catalog.pg_enum e ON t.oid = e.enumtypid
             WHERE t.typtype = 'e'
-            ORDER BY "name", e.oid -- e.oid gives us label order
         """
 
+        if schema != '*':
+            SQL_ENUMS += "AND n.nspname = :schema "
+
+        # e.oid gives us label order within an enum
+        SQL_ENUMS += 'ORDER BY "schema", "name", e.oid'
+
         s = sql.text(SQL_ENUMS, typemap={
             'attname': sqltypes.Unicode,
             'label': sqltypes.Unicode})
+
+        if schema != '*':
+            s = s.bindparams(schema=schema)
+
         c = connection.execute(s)
 
-        enums = {}
+        enums = []
+        enum_by_name = {}
         for enum in c.fetchall():
-            if enum['visible']:
-                # 'visible' just means whether or not the enum is in a
-                # schema that's on the search path -- or not overridden by
-                # a schema with higher precedence. If it's not visible,
-                # it will be prefixed with the schema-name when it's used.
-                name = enum['name']
-            else:
-                name = "%s.%s" % (enum['schema'], enum['name'])
-
-            if name in enums:
-                enums[name]['labels'].append(enum['label'])
+            key = (enum['schema'], enum['name'])
+            if key in enum_by_name:
+                enum_by_name[key]['labels'].append(enum['label'])
             else:
-                enums[name] = {
+                enum_by_name[key] = enum_rec = {
+                    'name': enum['name'],
+                    'schema': enum['schema'],
+                    'visible': enum['visible'],
                     'labels': [enum['label']],
                 }
+                enums.append(enum_rec)
 
         return enums