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

4 files changed, 101 insertions, 32 deletions

diff --git a/lib/sqlalchemy/dialects/mssql/base.py b/lib/sqlalchemy/dialects/mssql/base.py
index ed17fb863..ae852f264 100644
--- 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 84247d046..ffd926c46 100644
--- 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 4114137d4..304f4475f 100644
--- a/lib/sqlalchemy/testing/requirements.py
+++ b/lib/sqlalchemy/testing/requirements.py
@@ -595,6 +595,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