Merge "Handle PostgreSQL enums in remote schemas"

3 files changed, 79 insertions, 24 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index ce809db9f..d68ab8ef5 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1344,8 +1344,8 @@ class ENUM(sqltypes.NativeForEmulated, sqltypes.Enum):
                 pg_enums = ddl_runner.memo['_pg_enums']
             else:
                 pg_enums = ddl_runner.memo['_pg_enums'] = set()
-            present = self.name in pg_enums
-            pg_enums.add(self.name)
+            present = (self.schema, self.name) in pg_enums
+            pg_enums.add((self.schema, self.name))
             return present
         else:
             return False
@@ -2580,20 +2580,26 @@ class PGDialect(default.DefaultDialect):
                      )
         c = connection.execute(s, table_oid=table_oid)
         rows = c.fetchall()
+
+        # dictionary with (name, ) if default search path or (schema, name)
+        # as keys
         domains = self._load_domains(connection)
+
+        # dictionary with (name, ) if default search path or (schema, name)
+        # as keys
         enums = dict(
-            (
-                "%s.%s" % (rec['schema'], rec['name'])
-                if not rec['visible'] else rec['name'], rec) for rec in
-            self._load_enums(connection, schema='*')
+            ((rec['name'], ), rec)
+            if rec['visible'] else ((rec['schema'], rec['name']), rec)
+            for rec in self._load_enums(connection, schema='*')
         )
 
         # format columns
         columns = []
-        for name, format_type, default, notnull, attnum, table_oid, \
+
+        for name, format_type, default_, notnull, attnum, table_oid, \
                 comment in rows:
             column_info = self._get_column_info(
-                name, format_type, default, notnull, domains, enums,
+                name, format_type, default_, notnull, domains, enums,
                 schema, comment)
             columns.append(column_info)
         return columns
@@ -2602,7 +2608,8 @@ class PGDialect(default.DefaultDialect):
                          notnull, domains, enums, schema, comment):
         def _handle_array_type(attype):
             return (
-                attype.replace('[]', ''), # strip '[]' from integer[], etc.
+                # strip '[]' from integer[], etc.
+                re.sub(r'\[\]$', '', attype),
                 attype.endswith('[]'),
             )
 
@@ -2610,12 +2617,12 @@ class PGDialect(default.DefaultDialect):
         # with time zone, geometry(POLYGON), etc.
         attype = re.sub(r'\(.*\)', '', format_type)
 
-        # strip quotes from case sensitive enum names
-        attype = re.sub(r'^"|"$', '', attype)
-
         # strip '[]' from integer[], etc. and check if an array
         attype, is_array = _handle_array_type(attype)
 
+        # strip quotes from case sensitive enum or domain names
+        enum_or_domain_key = tuple(util.quoted_token_parser(attype))
+
         nullable = not notnull
 
         charlen = re.search(r'\(([\d,]+)\)', format_type)
@@ -2668,21 +2675,24 @@ class PGDialect(default.DefaultDialect):
             args = (int(charlen),)
 
         while True:
+            # looping here to suit nested domains
             if attype in self.ischema_names:
                 coltype = self.ischema_names[attype]
                 break
-            elif attype in enums:
-                enum = enums[attype]
+            elif enum_or_domain_key in enums:
+                enum = enums[enum_or_domain_key]
                 coltype = ENUM
                 kwargs['name'] = enum['name']
                 if not enum['visible']:
                     kwargs['schema'] = enum['schema']
                 args = tuple(enum['labels'])
                 break
-            elif attype in domains:
-                domain = domains[attype]
+            elif enum_or_domain_key in domains:
+                domain = domains[enum_or_domain_key]
                 attype = domain['attype']
                 attype, is_array = _handle_array_type(attype)
+                # strip quotes from case sensitive enum or domain names
+                enum_or_domain_key = tuple(util.quoted_token_parser(attype))
                 # A table can't override whether the domain is nullable.
                 nullable = domain['nullable']
                 if domain['default'] and not default:
@@ -3166,16 +3176,16 @@ class PGDialect(default.DefaultDialect):
         for domain in c.fetchall():
             # strip (30) from character varying(30)
             attype = re.search(r'([^\(]+)', domain['attype']).group(1)
+            # 'visible' just means whether or not the domain 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.
             if domain['visible']:
-                # 'visible' just means whether or not the domain 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 = domain['name']
+                key = (domain['name'], )
             else:
-                name = "%s.%s" % (domain['schema'], domain['name'])
+                key = (domain['schema'], domain['name'])
 
-            domains[name] = {
+            domains[key] = {
                 'attype': attype,
                 'nullable': domain['nullable'],
                 'default': domain['default']
diff --git a/lib/sqlalchemy/util/__init__.py b/lib/sqlalchemy/util/__init__.py
index 031376d78..9229d0797 100644
--- a/lib/sqlalchemy/util/__init__.py
+++ b/lib/sqlalchemy/util/__init__.py
@@ -34,7 +34,7 @@ from .langhelpers import iterate_attributes, class_hierarchy, \
     classproperty, set_creation_order, warn_exception, warn, NoneType,\
     constructor_copy, methods_equivalent, chop_traceback, asint,\
     generic_repr, counter, PluginLoader, hybridproperty, hybridmethod, \
-    safe_reraise,\
+    safe_reraise, quoted_token_parser,\
     get_callable_argspec, only_once, attrsetter, ellipses_string, \
     warn_limited, map_bits, MemoizedSlots, EnsureKWArgType, wrap_callable
 
diff --git a/lib/sqlalchemy/util/langhelpers.py b/lib/sqlalchemy/util/langhelpers.py
index 8815ed837..6a1db7a98 100644
--- a/lib/sqlalchemy/util/langhelpers.py
+++ b/lib/sqlalchemy/util/langhelpers.py
@@ -1422,3 +1422,48 @@ def wrap_callable(wrapper, fn):
             _f.__doc__ = fn.__doc__
 
         return _f
+
+
+def quoted_token_parser(value):
+    """Parse a dotted identifier with accomodation for quoted names.
+
+    Includes support for SQL-style double quotes as a literal character.
+
+    E.g.::
+
+        >>> quoted_token_parser("name")
+        ["name"]
+        >>> quoted_token_parser("schema.name")
+        ["schema", "name"]
+        >>> quoted_token_parser('"Schema"."Name"')
+        ['Schema', 'Name']
+        >>> quoted_token_parser('"Schema"."Name""Foo"')
+        ['Schema', 'Name""Foo']
+
+    """
+
+    if '"' not in value:
+        return value.split(".")
+
+    # 0 = outside of quotes
+    # 1 = inside of quotes
+    state = 0
+    result = [[]]
+    idx = 0
+    lv = len(value)
+    while idx < lv:
+        char = value[idx]
+        if char == '"':
+            if state == 1 and idx < lv - 1 and value[idx + 1] == '"':
+                result[-1].append('"')
+                idx += 1
+            else:
+                state ^= 1
+        elif char == "." and state == 0:
+            result.append([])
+        else:
+            result[-1].append(char)
+        idx += 1
+
+    return ["".join(token) for token in result]
+