Merge "Improve handling of covering indexes"

8 files changed, 213 insertions, 77 deletions

diff --git a/doc/build/changelog/unreleased_14/4458.rst b/doc/build/changelog/unreleased_14/4458.rst
new file mode 100644
index 000000000..976f51d52
--- /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 b398610b2..519d74d89 100644
--- a/lib/sqlalchemy/dialects/mssql/base.py
+++ b/lib/sqlalchemy/dialects/mssql/base.py
@@ -2873,6 +2873,7 @@ class MSDialect(default.DefaultDialect):
                 "name": row["name"],
                 "unique": row["is_unique"] == 1,
                 "column_names": [],
+                "include_columns": [],
             }
 
             if row["filter_definition"] is not None:
@@ -2882,7 +2883,8 @@ class MSDialect(default.DefaultDialect):
 
         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 "
@@ -2900,7 +2902,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 93ee7ca29..3ef87620f 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -668,6 +668,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
@@ -2235,8 +2249,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(
@@ -2248,12 +2273,10 @@ 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:
             whereclause = coercions.expect(
                 roles.DDLExpressionRole, whereclause
@@ -2263,6 +2286,7 @@ class PGDDLCompiler(compiler.DDLCompiler):
                 whereclause, include_table=False, literal_binds=True
             )
             text += " WHERE " + where_compiled
+
         return text
 
     def visit_drop_index(self, drop):
@@ -2732,6 +2756,7 @@ class PGDialect(default.DefaultDialect):
             schema.Index,
             {
                 "using": False,
+                "include": None,
                 "where": None,
                 "ops": {},
                 "concurrently": False,
@@ -3717,13 +3742,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.
@@ -3772,6 +3799,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 ed4c37933..97413d32b 100644
--- a/lib/sqlalchemy/testing/requirements.py
+++ b/lib/sqlalchemy/testing/requirements.py
@@ -599,6 +599,10 @@ class SuiteRequirements(Requirements):
         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."""
         return exclusions.open()
diff --git a/lib/sqlalchemy/testing/suite/test_reflection.py b/lib/sqlalchemy/testing/suite/test_reflection.py
index 94ec22c1e..3c10a45f6 100644
--- 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 30541ab06..2dd64d9bc 100644
--- a/test/dialect/postgresql/test_compiler.py
+++ b/test/dialect/postgresql/test_compiler.py
@@ -1734,6 +1734,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 2385a8f7a..263684e12 100644
--- a/test/dialect/postgresql/test_reflection.py
+++ b/test/dialect/postgresql/test_reflection.py
@@ -146,7 +146,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")
@@ -154,7 +162,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(
@@ -1074,7 +1092,11 @@ class ReflectionTest(AssertsCompiledSQL, 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")
@@ -1098,19 +1120,20 @@ class ReflectionTest(AssertsCompiledSQL, 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)
@@ -1136,17 +1159,17 @@ class ReflectionTest(AssertsCompiledSQL, 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_(
@@ -1176,14 +1199,17 @@ class ReflectionTest(AssertsCompiledSQL, 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
@@ -1542,18 +1568,19 @@ class ReflectionTest(AssertsCompiledSQL, 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 e0e9ce60c..28f82c3b8 100644
--- 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"])