Merge pull request #11 from malor/fix_uc_reflection

Fix unique constraints reflection in SQLite and PostgreSQL

3 files changed, 26 insertions, 13 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 16ace0583..9d89fe160 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -205,6 +205,7 @@ underlying CREATE INDEX command, so it *must* be a valid index type for your
 version of PostgreSQL.
 
 """
+from collections import defaultdict
 import re
 
 from ... import sql, schema, exc, util
@@ -2010,25 +2011,31 @@ class PGDialect(default.DefaultDialect):
         UNIQUE_SQL = """
             SELECT
                 cons.conname as name,
-                ARRAY_AGG(a.attname) as column_names
+                cons.conkey as key,
+                a.attnum as col_num,
+                a.attname as col_name
             FROM
                 pg_catalog.pg_constraint cons
-                left outer join pg_attribute a
-                    on cons.conrelid = a.attrelid and a.attnum = ANY(cons.conkey)
+                join pg_attribute a
+                  on cons.conrelid = a.attrelid AND a.attnum = ANY(cons.conkey)
             WHERE
                 cons.conrelid = :table_oid AND
                 cons.contype = 'u'
-            GROUP BY
-                cons.conname
         """
 
-        t = sql.text(UNIQUE_SQL,
-                     typemap={'column_names': ARRAY(sqltypes.Unicode)})
+        t = sql.text(UNIQUE_SQL, typemap={'col_name': sqltypes.Unicode})
         c = connection.execute(t, table_oid=table_oid)
 
+        uniques = defaultdict(lambda: defaultdict(dict))
+        for row in c.fetchall():
+            uc = uniques[row.name]
+            uc["key"] = row.key
+            uc["cols"][row.col_num] = row.col_name
+
         return [
-            {'name': row.name, 'column_names': row.column_names}
-            for row in c.fetchall()
+            {'name': name,
+             'column_names': [uc["cols"][i] for i in uc["key"]]}
+            for name, uc in uniques.items()
         ]
 
     def _load_enums(self, connection):
diff --git a/lib/sqlalchemy/dialects/sqlite/base.py b/lib/sqlalchemy/dialects/sqlite/base.py
index 3e2a158a0..787fdec17 100644
--- a/lib/sqlalchemy/dialects/sqlite/base.py
+++ b/lib/sqlalchemy/dialects/sqlite/base.py
@@ -933,7 +933,8 @@ class SQLiteDialect(default.DefaultDialect):
 
         UNIQUE_PATTERN = 'CONSTRAINT (\w+) UNIQUE \(([^\)]+)\)'
         return [
-            {'name': name, 'column_names': [c.strip() for c in cols.split(',')]}
+            {'name': name,
+             'column_names': [col.strip(' "') for col in cols.split(',')]}
             for name, cols in re.findall(UNIQUE_PATTERN, table_data)
         ]
 
diff --git a/lib/sqlalchemy/testing/suite/test_reflection.py b/lib/sqlalchemy/testing/suite/test_reflection.py
index 858ca8a77..255fd5595 100644
--- a/lib/sqlalchemy/testing/suite/test_reflection.py
+++ b/lib/sqlalchemy/testing/suite/test_reflection.py
@@ -370,9 +370,10 @@ class ComponentReflectionTest(fixtures.TablesTest):
     def _test_get_unique_constraints(self, schema=None):
         uniques = sorted(
             [
+                {'name': 'unique_a', 'column_names': ['a']},
                 {'name': 'unique_a_b_c', 'column_names': ['a', 'b', 'c']},
-                {'name': 'unique_a_c', 'column_names': ['a', 'c']},
-                {'name': 'unique_b_c', 'column_names': ['b', 'c']},
+                {'name': 'unique_c_a_b', 'column_names': ['c', 'a', 'b']},
+                {'name': 'unique_asc_key', 'column_names': ['asc', 'key']},
             ],
             key=operator.itemgetter('name')
         )
@@ -382,6 +383,9 @@ class ComponentReflectionTest(fixtures.TablesTest):
             Column('a', sa.String(20)),
             Column('b', sa.String(30)),
             Column('c', sa.Integer),
+            # reserved identifiers
+            Column('asc', sa.String(30)),
+            Column('key', sa.String(30)),
             schema=schema
         )
         for uc in uniques:
@@ -396,7 +400,8 @@ class ComponentReflectionTest(fixtures.TablesTest):
             key=operator.itemgetter('name')
         )
 
-        eq_(uniques, reflected)
+        for orig, refl in zip(uniques, reflected):
+            eq_(orig, refl)
 
 
     @testing.provide_metadata