Improve handling of covering indexes

Improved support for covering indexes (with INCLUDE columns). Added the
ability for postgresql to render CREATE INDEX statements with an INCLUDE
clause from Core. Index reflection also report INCLUDE columns separately
for both mssql and postgresql (11+).

Fixes: #4458
Change-Id: If0b82103fbc898cdaeaf6a6d2d421c732744acd6

diff --git a/doc/build/changelog/unreleased_14/4458.rst b/doc/build/changelog/unreleased_14/4458.rst
new file mode 100644 (file)
index 0000000..976f51d
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/4458.rst
@@ -0,0 +1,8 @@
+.. change::
+    :tags: mssql, postgresql, reflection, schema, usecase
+    :tickets: 4458
+
+    Improved support for covering indexes (with INCLUDE columns). Added the
+    ability for postgresql to render CREATE INDEX statements with an INCLUDE
+    clause from Core. Index reflection also report INCLUDE columns separately
+    for both mssql and postgresql (11+).

diff --git a/lib/sqlalchemy/dialects/mssql/base.py b/lib/sqlalchemy/dialects/mssql/base.py
index ed17fb8631e5a4a96c358c32bb9babf819a46b02..ae852f26431f4dd1120bc5096bc8f410259ad893 100644 (file)
--- a/lib/sqlalchemy/dialects/mssql/base.py
+++ b/lib/sqlalchemy/dialects/mssql/base.py
@@ -2864,10 +2864,12 @@ class MSDialect(default.DefaultDialect):
                 "name": row["name"],
                 "unique": row["is_unique"] == 1,
                 "column_names": [],
+                "include_columns": [],
             }
         rp = connection.execution_options(future_result=True).execute(
             sql.text(
-                "select ind_col.index_id, ind_col.object_id, col.name "
+                "select ind_col.index_id, ind_col.object_id, col.name, "
+                "ind_col.is_included_column "
                 "from sys.columns as col "
                 "join sys.tables as tab on tab.object_id=col.object_id "
                 "join sys.index_columns as ind_col on "
@@ -2885,7 +2887,14 @@ class MSDialect(default.DefaultDialect):
         )
         for row in rp.mappings():
             if row["index_id"] in indexes:
-                indexes[row["index_id"]]["column_names"].append(row["name"])
+                if row["is_included_column"]:
+                    indexes[row["index_id"]]["include_columns"].append(
+                        row["name"]
+                    )
+                else:
+                    indexes[row["index_id"]]["column_names"].append(
+                        row["name"]
+                    )
 
         return list(indexes.values())
 

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 84247d0467e45a4f049e6cdfadd48618a27c32d7..ffd926c460b417603af0d09ac1a5077a414f2c9d 100644 (file)
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -670,6 +670,20 @@ PostgreSQL-Specific Index Options
 Several extensions to the :class:`.Index` construct are available, specific
 to the PostgreSQL dialect.
 
+Covering Indexes
+^^^^^^^^^^^^^^^^
+
+The ``postgresql_include`` option renders INCLUDE(colname) for the given
+string names::
+
+    Index("my_index", table.c.x, postgresql_include=['y'])
+
+would render the index as ``CREATE INDEX my_index ON table (x) INCLUDE (y)``
+
+Note that this feature requires PostgreSQL 11 or later.
+
+.. versionadded:: 1.4
+
 .. _postgresql_partial_indexes:
 
 Partial Indexes
@@ -2237,8 +2251,19 @@ class PGDDLCompiler(compiler.DDLCompiler):
             )
         )
 
-        withclause = index.dialect_options["postgresql"]["with"]
+        includeclause = index.dialect_options["postgresql"]["include"]
+        if includeclause:
+            inclusions = [
+                index.table.c[col]
+                if isinstance(col, util.string_types)
+                else col
+                for col in includeclause
+            ]
+            text += " INCLUDE (%s)" % ", ".join(
+                [preparer.quote(c.name) for c in inclusions]
+            )
 
+        withclause = index.dialect_options["postgresql"]["with"]
         if withclause:
             text += " WITH (%s)" % (
                 ", ".join(
@@ -2250,17 +2275,16 @@ class PGDDLCompiler(compiler.DDLCompiler):
             )
 
         tablespace_name = index.dialect_options["postgresql"]["tablespace"]
-
         if tablespace_name:
             text += " TABLESPACE %s" % preparer.quote(tablespace_name)
 
         whereclause = index.dialect_options["postgresql"]["where"]
-
         if whereclause is not None:
             where_compiled = self.sql_compiler.process(
                 whereclause, include_table=False, literal_binds=True
             )
             text += " WHERE " + where_compiled
+
         return text
 
     def visit_drop_index(self, drop):
@@ -2731,6 +2755,7 @@ class PGDialect(default.DefaultDialect):
             schema.Index,
             {
                 "using": False,
+                "include": None,
                 "where": None,
                 "ops": {},
                 "concurrently": False,
@@ -3722,13 +3747,15 @@ class PGDialect(default.DefaultDialect):
                 # included columns, which are merely stored and do not
                 # participate in the index semantics"
                 if indnkeyatts and idx_keys[indnkeyatts:]:
-                    util.warn(
-                        "INCLUDE columns for covering index %s "
-                        "ignored during reflection" % (idx_name,)
-                    )
+                    # this is a "covering index" which has INCLUDE columns
+                    # as well as regular index columns
+                    inc_keys = idx_keys[indnkeyatts:]
                     idx_keys = idx_keys[:indnkeyatts]
+                else:
+                    inc_keys = []
 
                 index["key"] = [int(k.strip()) for k in idx_keys]
+                index["inc"] = [int(k.strip()) for k in inc_keys]
 
                 # (new in pg 8.3)
                 # "pg_index.indoption" is list of ints, one per column/expr.
@@ -3774,6 +3801,8 @@ class PGDialect(default.DefaultDialect):
                 "unique": idx["unique"],
                 "column_names": [idx["cols"][i] for i in idx["key"]],
             }
+            if self.server_version_info >= (11, 0):
+                entry["include_columns"] = [idx["cols"][i] for i in idx["inc"]]
             if "duplicates_constraint" in idx:
                 entry["duplicates_constraint"] = idx["duplicates_constraint"]
             if "sorting" in idx:

diff --git a/lib/sqlalchemy/testing/requirements.py b/lib/sqlalchemy/testing/requirements.py
index 4114137d4e9608c4542ceb57a36348d17e725cee..304f4475f381ef0d4d70fb5bd9d8d1e7ca779b8e 100644 (file)
--- a/lib/sqlalchemy/testing/requirements.py
+++ b/lib/sqlalchemy/testing/requirements.py
@@ -594,6 +594,10 @@ class SuiteRequirements(Requirements):
     def index_reflection(self):
         return exclusions.open()
 
+    @property
+    def index_reflects_included_columns(self):
+        return exclusions.closed()
+
     @property
     def indexes_with_ascdesc(self):
         """target database supports CREATE INDEX with per-column ASC/DESC."""

diff --git a/lib/sqlalchemy/testing/suite/test_reflection.py b/lib/sqlalchemy/testing/suite/test_reflection.py
index 94ec22c1e6bac3e3feb59601ac49af1a4f335e10..3c10a45f62f274f07534c94a0b1482e26f8426d7 100644 (file)
--- a/lib/sqlalchemy/testing/suite/test_reflection.py
+++ b/lib/sqlalchemy/testing/suite/test_reflection.py
@@ -2,6 +2,7 @@ import operator
 import re
 
 import sqlalchemy as sa
+from sqlalchemy import func
 from .. import config
 from .. import engines
 from .. import eq_
@@ -1013,32 +1014,62 @@ class ComponentReflectionTest(fixtures.TablesTest):
     @testing.requires.indexes_with_expressions
     @testing.provide_metadata
     def test_reflect_expression_based_indexes(self):
-        Table(
+        t = Table(
             "t",
             self.metadata,
             Column("x", String(30)),
             Column("y", String(30)),
         )
-        event.listen(
-            self.metadata,
-            "after_create",
-            DDL("CREATE INDEX t_idx ON t(lower(x), lower(y))"),
-        )
-        event.listen(
-            self.metadata, "after_create", DDL("CREATE INDEX t_idx_2 ON t(x)")
-        )
-        self.metadata.create_all()
 
-        insp = inspect(self.metadata.bind)
+        Index("t_idx", func.lower(t.c.x), func.lower(t.c.y))
+
+        Index("t_idx_2", t.c.x)
+
+        self.metadata.create_all(testing.db)
+
+        insp = inspect(testing.db)
+
+        expected = [
+            {"name": "t_idx_2", "column_names": ["x"], "unique": False}
+        ]
+        if testing.requires.index_reflects_included_columns.enabled:
+            expected[0]["include_columns"] = []
 
         with expect_warnings(
             "Skipped unsupported reflection of expression-based index t_idx"
         ):
             eq_(
-                insp.get_indexes("t"),
-                [{"name": "t_idx_2", "column_names": ["x"], "unique": 0}],
+                insp.get_indexes("t"), expected,
             )
 
+    @testing.requires.index_reflects_included_columns
+    @testing.provide_metadata
+    def test_reflect_covering_index(self):
+        t = Table(
+            "t",
+            self.metadata,
+            Column("x", String(30)),
+            Column("y", String(30)),
+        )
+        idx = Index("t_idx", t.c.x)
+        idx.dialect_options[testing.db.name]["include"] = ["y"]
+
+        self.metadata.create_all(testing.db)
+
+        insp = inspect(testing.db)
+
+        eq_(
+            insp.get_indexes("t"),
+            [
+                {
+                    "name": "t_idx",
+                    "column_names": ["x"],
+                    "include_columns": ["y"],
+                    "unique": False,
+                }
+            ],
+        )
+
     @testing.requires.unique_constraint_reflection
     def test_get_unique_constraints(self):
         self._test_get_unique_constraints()
@@ -1061,17 +1092,13 @@ class ComponentReflectionTest(fixtures.TablesTest):
         indexes = insp.get_indexes(table_name)
         for ind in indexes:
             ind.pop("dialect_options", None)
+        expected = [
+            {"unique": False, "column_names": ["foo"], "name": "user_tmp_ix"}
+        ]
+        if testing.requires.index_reflects_included_columns.enabled:
+            expected[0]["include_columns"] = []
         eq_(
-            # TODO: we need to add better filtering for indexes/uq constraints
-            # that are doubled up
-            [idx for idx in indexes if idx["name"] == "user_tmp_ix"],
-            [
-                {
-                    "unique": False,
-                    "column_names": ["foo"],
-                    "name": "user_tmp_ix",
-                }
-            ],
+            [idx for idx in indexes if idx["name"] == "user_tmp_ix"], expected,
         )
 
     @testing.requires.unique_constraint_reflection

diff --git a/test/dialect/postgresql/test_compiler.py b/test/dialect/postgresql/test_compiler.py
index 556601fc62ba591977a63e7453f4d6198cb49295..c20498b2dff46c57801d5e39c8c8ca55bfd12b50 100644 (file)
--- a/test/dialect/postgresql/test_compiler.py
+++ b/test/dialect/postgresql/test_compiler.py
@@ -1721,6 +1721,34 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
             "(INCREMENT BY 7 START WITH 4))",
         )
 
+    def test_index_extra_include_1(self):
+        metadata = MetaData()
+        tbl = Table(
+            "test",
+            metadata,
+            Column("x", Integer),
+            Column("y", Integer),
+            Column("z", Integer),
+        )
+        idx = Index("foo", tbl.c.x, postgresql_include=["y"])
+        self.assert_compile(
+            schema.CreateIndex(idx), "CREATE INDEX foo ON test (x) INCLUDE (y)"
+        )
+
+    def test_index_extra_include_2(self):
+        metadata = MetaData()
+        tbl = Table(
+            "test",
+            metadata,
+            Column("x", Integer),
+            Column("y", Integer),
+            Column("z", Integer),
+        )
+        idx = Index("foo", tbl.c.x, postgresql_include=[tbl.c.y])
+        self.assert_compile(
+            schema.CreateIndex(idx), "CREATE INDEX foo ON test (x) INCLUDE (y)"
+        )
+
 
 class InsertOnConflictTest(fixtures.TestBase, AssertsCompiledSQL):
     __dialect__ = postgresql.dialect()

diff --git a/test/dialect/postgresql/test_reflection.py b/test/dialect/postgresql/test_reflection.py
index ec9328c2fb1ea45a9fd2809329cfdbf9dd8bbb5c..e088cad01d9169864eed51ff7caef07d28ea8af5 100644 (file)
--- a/test/dialect/postgresql/test_reflection.py
+++ b/test/dialect/postgresql/test_reflection.py
@@ -144,7 +144,15 @@ class PartitionedReflectionTest(fixtures.TablesTest, AssertsExecutionResults):
     def test_reflect_index(self):
         idx = inspect(testing.db).get_indexes("data_values")
         eq_(
-            idx, [{"column_names": ["q"], "name": "my_index", "unique": False}]
+            idx,
+            [
+                {
+                    "name": "my_index",
+                    "unique": False,
+                    "column_names": ["q"],
+                    "include_columns": [],
+                }
+            ],
         )
 
     @testing.only_on("postgresql >= 11")
@@ -152,7 +160,17 @@ class PartitionedReflectionTest(fixtures.TablesTest, AssertsExecutionResults):
         idx = inspect(testing.db).get_indexes("data_values_4_10")
         # note the name appears to be generated by PG, currently
         # 'data_values_4_10_q_idx'
-        eq_(idx, [{"column_names": ["q"], "name": mock.ANY, "unique": False}])
+        eq_(
+            idx,
+            [
+                {
+                    "column_names": ["q"],
+                    "include_columns": [],
+                    "name": mock.ANY,
+                    "unique": False,
+                }
+            ],
+        )
 
 
 class MaterializedViewReflectionTest(
@@ -1031,7 +1049,11 @@ class ReflectionTest(fixtures.TestBase):
         conn.exec_driver_sql("ALTER TABLE t RENAME COLUMN x to y")
 
         ind = testing.db.dialect.get_indexes(conn, "t", None)
-        eq_(ind, [{"unique": False, "column_names": ["y"], "name": "idx1"}])
+        expected = [{"name": "idx1", "unique": False, "column_names": ["y"]}]
+        if testing.requires.index_reflects_included_columns.enabled:
+            expected[0]["include_columns"] = []
+
+        eq_(ind, expected)
         conn.close()
 
     @testing.fails_if("postgresql < 8.2", "reloptions not supported")
@@ -1055,19 +1077,20 @@ class ReflectionTest(fixtures.TestBase):
             )
 
             ind = testing.db.dialect.get_indexes(conn, "t", None)
-            eq_(
-                ind,
-                [
-                    {
-                        "unique": False,
-                        "column_names": ["x"],
-                        "name": "idx1",
-                        "dialect_options": {
-                            "postgresql_with": {"fillfactor": "50"}
-                        },
-                    }
-                ],
-            )
+
+            expected = [
+                {
+                    "unique": False,
+                    "column_names": ["x"],
+                    "name": "idx1",
+                    "dialect_options": {
+                        "postgresql_with": {"fillfactor": "50"}
+                    },
+                }
+            ]
+            if testing.requires.index_reflects_included_columns.enabled:
+                expected[0]["include_columns"] = []
+            eq_(ind, expected)
 
             m = MetaData()
             t1 = Table("t", m, autoload_with=conn)
@@ -1093,17 +1116,17 @@ class ReflectionTest(fixtures.TestBase):
             conn.exec_driver_sql("CREATE INDEX idx1 ON t USING gin (x)")
 
             ind = testing.db.dialect.get_indexes(conn, "t", None)
-            eq_(
-                ind,
-                [
-                    {
-                        "unique": False,
-                        "column_names": ["x"],
-                        "name": "idx1",
-                        "dialect_options": {"postgresql_using": "gin"},
-                    }
-                ],
-            )
+            expected = [
+                {
+                    "unique": False,
+                    "column_names": ["x"],
+                    "name": "idx1",
+                    "dialect_options": {"postgresql_using": "gin"},
+                }
+            ]
+            if testing.requires.index_reflects_included_columns.enabled:
+                expected[0]["include_columns"] = []
+            eq_(ind, expected)
             m = MetaData()
             t1 = Table("t", m, autoload_with=conn)
             eq_(
@@ -1133,14 +1156,17 @@ class ReflectionTest(fixtures.TestBase):
             # [{'column_names': ['x', 'name'],
             #  'name': 'idx1', 'unique': False}]
 
-            with testing.expect_warnings(
-                "INCLUDE columns for "
-                "covering index idx1 ignored during reflection"
-            ):
-                ind = testing.db.dialect.get_indexes(conn, "t", None)
+            ind = testing.db.dialect.get_indexes(conn, "t", None)
             eq_(
                 ind,
-                [{"unique": False, "column_names": ["x"], "name": "idx1"}],
+                [
+                    {
+                        "unique": False,
+                        "column_names": ["x"],
+                        "include_columns": ["name"],
+                        "name": "idx1",
+                    }
+                ],
             )
 
     @testing.provide_metadata
@@ -1499,18 +1525,19 @@ class ReflectionTest(fixtures.TestBase):
 
         # PostgreSQL will create an implicit index for an exclude constraint.
         # we don't reflect the EXCLUDE yet.
-        eq_(
-            insp.get_indexes("t"),
-            [
-                {
-                    "unique": False,
-                    "name": "quarters_period_excl",
-                    "duplicates_constraint": "quarters_period_excl",
-                    "dialect_options": {"postgresql_using": "gist"},
-                    "column_names": ["period"],
-                }
-            ],
-        )
+        expected = [
+            {
+                "unique": False,
+                "name": "quarters_period_excl",
+                "duplicates_constraint": "quarters_period_excl",
+                "dialect_options": {"postgresql_using": "gist"},
+                "column_names": ["period"],
+            }
+        ]
+        if testing.requires.index_reflects_included_columns.enabled:
+            expected[0]["include_columns"] = []
+
+        eq_(insp.get_indexes("t"), expected)
 
         # reflection corrects for the dupe
         reflected = Table("t", MetaData(testing.db), autoload=True)

diff --git a/test/requirements.py b/test/requirements.py
index 145d87d753c4342cd973eadf0bec85adde666e73..4655daf7aac5adacc81d56fc2e8e2adc9f8248a5 100644 (file)
--- a/test/requirements.py
+++ b/test/requirements.py
@@ -1688,3 +1688,7 @@ class DefaultRequirements(SuiteRequirements):
     @property
     def identity_columns_standard(self):
         return self.identity_columns + skip_if("mssql")
+
+    @property
+    def index_reflects_included_columns(self):
+        return only_on(["postgresql >= 11", "mssql"])