Merge "Improve handling of covering indexes"

1 files changed, 36 insertions, 7 deletions

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: