add reflection for expression-based index and partial index of postgresqlpr/242

4 files changed, 97 insertions, 155 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index b16a82e04..855b77fb6 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1356,6 +1356,7 @@ class PGDDLCompiler(compiler.DDLCompiler):
 
 
 class PGTypeCompiler(compiler.GenericTypeCompiler):
+
     def visit_TSVECTOR(self, type, **kw):
         return "TSVECTOR"
 
@@ -2329,110 +2330,64 @@ class PGDialect(default.DefaultDialect):
         table_oid = self.get_table_oid(connection, table_name, schema,
                                        info_cache=kw.get('info_cache'))
 
-        # cast indkey as varchar since it's an int2vector,
-        # returned as a list by some drivers such as pypostgresql
-
-        if self.server_version_info < (8, 5):
-            IDX_SQL = """
-              SELECT
-                  i.relname as relname,
-                  ix.indisunique, ix.indexprs, ix.indpred,
-                  a.attname, a.attnum, NULL, ix.indkey%s,
-                  %s, am.amname
-              FROM
-                  pg_class t
-                        join pg_index ix on t.oid = ix.indrelid
-                        join pg_class i on i.oid = ix.indexrelid
-                        left outer join
-                            pg_attribute a
-                            on t.oid = a.attrelid and %s
-                        left outer join
-                            pg_am am
-                            on i.relam = am.oid
-              WHERE
-                  t.relkind IN ('r', 'v', 'f', 'm')
-                  and t.oid = :table_oid
-                  and ix.indisprimary = 'f'
-              ORDER BY
-                  t.relname,
-                  i.relname
-            """ % (
-                # version 8.3 here was based on observing the
-                # cast does not work in PG 8.2.4, does work in 8.3.0.
-                # nothing in PG changelogs regarding this.
-                "::varchar" if self.server_version_info >= (8, 3) else "",
-                "i.reloptions" if self.server_version_info >= (8, 2)
-                else "NULL",
-                self._pg_index_any("a.attnum", "ix.indkey")
-            )
-        else:
-            IDX_SQL = """
-              SELECT
-                  i.relname as relname,
-                  ix.indisunique, ix.indexprs, ix.indpred,
-                  a.attname, a.attnum, c.conrelid, ix.indkey::varchar,
-                  i.reloptions, am.amname
-              FROM
-                  pg_class t
-                        join pg_index ix on t.oid = ix.indrelid
-                        join pg_class i on i.oid = ix.indexrelid
-                        left outer join
-                            pg_attribute a
-                            on t.oid = a.attrelid and a.attnum = ANY(ix.indkey)
-                        left outer join
-                            pg_constraint c
-                            on (ix.indrelid = c.conrelid and
-                                ix.indexrelid = c.conindid and
-                                c.contype in ('p', 'u', 'x'))
-                        left outer join
-                            pg_am am
-                            on i.relam = am.oid
-              WHERE
-                  t.relkind IN ('r', 'v', 'f', 'm')
-                  and t.oid = :table_oid
-                  and ix.indisprimary = 'f'
-              ORDER BY
-                  t.relname,
-                  i.relname
-            """
+        IDX_SQL = """
+          SELECT
+              i.relname as relname,
+              ix.indisunique,
+              pg_get_indexdef(i.oid, s.i + 1, false),
+              s.i,
+              pg_get_expr(ix.indpred, ix.indrelid),
+              c.conrelid,
+              %s,
+              am.amname
+          FROM
+              pg_class t
+                    join pg_index ix on t.oid = ix.indrelid
+                    join pg_class i on i.oid = ix.indexrelid
+                    join generate_series(0, current_setting('max_index_keys')::int - 1) s(i)
+                        on ix.indkey[s.i] is NOT NULL
+                    left outer join
+                        pg_constraint c
+                        on (ix.indrelid = c.conrelid and
+                        ix.indexrelid = c.conindid and
+                        c.contype in ('p', 'u', 'x'))
+                    left outer join
+                        pg_am am
+                        on i.relam = am.oid
+          WHERE
+              t.relkind IN ('r', 'v', 'f', 'm')
+              and t.oid = :table_oid
+              and ix.indisprimary = 'f'
+          ORDER BY
+              t.relname,
+              i.relname,
+              s.i
+        """ % (
+            "i.reloptions" if self.server_version_info >= (8, 2) else "NULL"
+        )
 
         t = sql.text(IDX_SQL, typemap={'attname': sqltypes.Unicode})
         c = connection.execute(t, table_oid=table_oid)
 
-        indexes = defaultdict(lambda: defaultdict(dict))
+        indexes = defaultdict(lambda: defaultdict(list))
 
         sv_idx_name = None
         for row in c.fetchall():
-            (idx_name, unique, expr, prd, col,
-             col_num, conrelid, idx_key, options, amname) = row
-
-            if expr:
-                if idx_name != sv_idx_name:
-                    util.warn(
-                        "Skipped unsupported reflection of "
-                        "expression-based index %s"
-                        % idx_name)
-                sv_idx_name = idx_name
-                continue
-
-            if prd and not idx_name == sv_idx_name:
-                util.warn(
-                    "Predicate of partial index %s ignored during reflection"
-                    % idx_name)
-                sv_idx_name = idx_name
+            (idx_name, unique, col, col_num, prd, conrelid, options, amname) = row
 
             has_idx = idx_name in indexes
             index = indexes[idx_name]
             if col is not None:
-                index['cols'][col_num] = col
+                index['cols'].append(col.strip('"'))
             if not has_idx:
-                index['key'] = [int(k.strip()) for k in idx_key.split()]
                 index['unique'] = unique
                 if conrelid is not None:
                     index['duplicates_constraint'] = idx_name
                 if options:
                     index['options'] = dict(
                         [option.split("=") for option in options])
+                if prd:
+                    index['where'] = prd
 
                 # it *might* be nice to include that this is 'btree' in the
                 # reflection info.  But we don't want an Index object
@@ -2446,16 +2401,19 @@ class PGDialect(default.DefaultDialect):
             entry = {
                 'name': name,
                 'unique': idx['unique'],
-                'column_names': [idx['cols'][i] for i in idx['key']]
+                'column_names': idx['cols']
             }
             if 'duplicates_constraint' in idx:
                 entry['duplicates_constraint'] = idx['duplicates_constraint']
             if 'options' in idx:
                 entry.setdefault(
-                    'dialect_options', {})["postgresql_with"] = idx['options']
+                    'dialect_options', {})['postgresql_with'] = idx['options']
             if 'amname' in idx:
                 entry.setdefault(
-                    'dialect_options', {})["postgresql_using"] = idx['amname']
+                    'dialect_options', {})['postgresql_using'] = idx['amname']
+            if 'where' in idx:
+                entry.setdefault(
+                    'dialect_options', {})['postgresql_where'] = idx['where']
             result.append(entry)
         return result
 
diff --git a/lib/sqlalchemy/engine/reflection.py b/lib/sqlalchemy/engine/reflection.py
index eaa5e2e48..833e53014 100644
--- a/lib/sqlalchemy/engine/reflection.py
+++ b/lib/sqlalchemy/engine/reflection.py
@@ -57,6 +57,7 @@ def cache(fn, self, con, *args, **kw):
 
 
 class Inspector(object):
+
     """Performs database schema inspection.
 
     The Inspector acts as a proxy to the reflection methods of the
@@ -727,26 +728,25 @@ class Inspector(object):
                 continue
             if duplicates:
                 continue
-            # look for columns by orig name in cols_by_orig_name,
-            # but support columns that are in-Python only as fallback
+
             idx_cols = []
             for c in columns:
-                try:
-                    idx_col = cols_by_orig_name[c] \
-                        if c in cols_by_orig_name else table.c[c]
-                except KeyError:
-                    util.warn(
-                        "%s key '%s' was not located in "
-                        "columns for table '%s'" % (
-                            flavor, c, table_name
-                        ))
+                if c in cols_by_orig_name:
+                    idx_col = cols_by_orig_name[c]
+                elif c in table.c:
+                    idx_col = table.c[c]
                 else:
-                    idx_cols.append(idx_col)
+                    # expression-based index
+                    idx_col = sql.text(c)
+
+                idx_cols.append(idx_col)
 
-            sa_schema.Index(
+            idx = sa_schema.Index(
                 name, *idx_cols,
                 **dict(list(dialect_options.items()) + [('unique', unique)])
             )
+            if idx.table is None:
+                table.append_constraint(idx)
 
     def _reflect_unique_constraints(
         self, table_name, schema, table, cols_by_orig_name,
diff --git a/test/dialect/postgresql/test_reflection.py b/test/dialect/postgresql/test_reflection.py
index 8da18108f..27874b28a 100644
--- a/test/dialect/postgresql/test_reflection.py
+++ b/test/dialect/postgresql/test_reflection.py
@@ -17,6 +17,7 @@ import re
 
 
 class ForeignTableReflectionTest(fixtures.TablesTest, AssertsExecutionResults):
+
     """Test reflection on foreign tables"""
 
     __requires__ = 'postgresql_test_dblink',
@@ -74,6 +75,7 @@ class ForeignTableReflectionTest(fixtures.TablesTest, AssertsExecutionResults):
 
 class MaterializedViewReflectionTest(
         fixtures.TablesTest, AssertsExecutionResults):
+
     """Test reflection on materialized views"""
 
     __only_on__ = 'postgresql >= 9.3'
@@ -142,6 +144,7 @@ class MaterializedViewReflectionTest(
 
 
 class DomainReflectionTest(fixtures.TestBase, AssertsExecutionResults):
+
     """Test PostgreSQL domains"""
 
     __only_on__ = 'postgresql > 8.3'
@@ -610,7 +613,7 @@ class ReflectionTest(fixtures.TestBase):
 
     @testing.provide_metadata
     def test_index_reflection(self):
-        """ Reflecting partial & expression-based indexes should warn
+        """ reflect partial & expression-based indexes
         """
 
         metadata = self.metadata
@@ -624,41 +627,43 @@ class ReflectionTest(fixtures.TestBase):
             Column(
                 'aname', String(20)))
         metadata.create_all()
-        testing.db.execute("""
-          create index idx1 on party ((id || name))
-        """)
-        testing.db.execute("""
-          create unique index idx2 on party (id) where name = 'test'
-        """)
-        testing.db.execute("""
-            create index idx3 on party using btree
-                (lower(name::text), lower(aname::text))
-        """)
-
-        def go():
-            m2 = MetaData(testing.db)
-            t2 = Table('party', m2, autoload=True)
-            assert len(t2.indexes) == 2
+        with testing.db.connect().execution_options(autocommit=True) as conn:
+            conn.execute("""
+                CREATE INDEX idx1 ON party (((id)::text || (name)::text))
+            """)
+            conn.execute("""
+                CREATE UNIQUE INDEX idx2 ON party (id) WHERE (name)::text = 'test'::text
+            """)
+            conn.execute("""
+                CREATE INDEX idx3 ON party USING btree
+                    (lower((name)::text), lower((aname)::text))
+            """)
 
-            # Make sure indexes are in the order we expect them in
+            m = MetaData()
+            t2 = Table('party', m, autoload_with=conn)
+            eq_(len(t2.indexes), 4)
 
             tmp = [(idx.name, idx) for idx in t2.indexes]
             tmp.sort()
-            r1, r2 = [idx[1] for idx in tmp]
-            assert r1.name == 'idx2'
-            assert r1.unique == True
-            assert r2.unique == False
-            assert [t2.c.id] == r1.columns
-            assert [t2.c.name] == r2.columns
-
-        testing.assert_warnings(
-            go,
-            ['Skipped unsupported reflection of '
-             'expression-based index idx1',
-             'Predicate of partial index idx2 ignored during '
-             'reflection',
-             'Skipped unsupported reflection of '
-             'expression-based index idx3'])
+            r1, r2, r3, r4 = [idx[1] for idx in tmp]
+
+            eq_(r1.unique, False)
+            eq_(r1.name, 'idx1')
+            eq_(r1.expressions[0].text, '(((id)::text || (name)::text))')
+
+            eq_(r2.unique, True)
+            eq_(r2.name, 'idx2')
+            eq_(r2.unique, True)
+            eq_(r2.columns, [t2.c.id])
+            eq_(r2.dialect_options['postgresql']['where'], "((name)::text = 'test'::text)")
+
+            eq_(r3.unique, False)
+            eq_(r3.name, 'idx3')
+            eq_(r3.expressions[0].text, 'lower((name)::text)')
+            eq_(r3.expressions[1].text, 'lower((aname)::text)')
+
+            eq_(r4.unique, False)
+            eq_(r4.columns, [t2.c.name])
 
     @testing.provide_metadata
     def test_index_reflection_modified(self):
diff --git a/test/engine/test_reflection.py b/test/engine/test_reflection.py
index f9799fda0..7ce71d06d 100644
--- a/test/engine/test_reflection.py
+++ b/test/engine/test_reflection.py
@@ -981,27 +981,6 @@ class ReflectionTest(fixtures.TestBase, ComparesTables):
         assert set([t2.c.name, t2.c.id]) == set(r2.columns)
         assert set([t2.c.name]) == set(r3.columns)
 
-    @testing.provide_metadata
-    def test_index_reflection_cols_busted(self):
-        t = Table('x', self.metadata,
-                  Column('a', Integer), Column('b', Integer))
-        sa.Index('x_ix', t.c.a, t.c.b)
-        self.metadata.create_all()
-
-        def mock_get_columns(self, connection, table_name, **kw):
-            return [
-                {"name": "b", "type": Integer, "primary_key": False}
-            ]
-
-        with testing.mock.patch.object(
-                testing.db.dialect, "get_columns", mock_get_columns):
-            m = MetaData()
-            with testing.expect_warnings(
-                    "index key 'a' was not located in columns"):
-                t = Table('x', m, autoload=True, autoload_with=testing.db)
-
-        eq_(list(t.indexes)[0].columns, [t.c.b])
-
     @testing.requires.views
     @testing.provide_metadata
     def test_views(self):