From 50830b2fdb712552b4a5e5292e9ad6448d60cd91 Mon Sep 17 00:00:00 2001
From: Marat Sharafutdinov <decaz89@gmail.com>
Date: Thu, 19 Mar 2020 00:19:17 +0300
Subject: [PATCH] Fixed reflection of covering indexes

---
 lib/sqlalchemy/dialects/postgresql/base.py | 18 ++++++++++++----
 test/dialect/postgresql/test_reflection.py | 24 ++++++++++++++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 8105244fa4..6348962af1 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -3259,7 +3259,8 @@ class PGDialect(default.DefaultDialect):
                   i.relname as relname,
                   ix.indisunique, ix.indexprs, ix.indpred,
                   a.attname, a.attnum, NULL, ix.indkey%s,
-                  %s, %s, am.amname
+                  %s, %s, am.amname,
+                  NULL as indnkeyatts
               FROM
                   pg_class t
                         join pg_index ix on t.oid = ix.indrelid
@@ -3296,7 +3297,8 @@ class PGDialect(default.DefaultDialect):
                   i.relname as relname,
                   ix.indisunique, ix.indexprs, ix.indpred,
                   a.attname, a.attnum, c.conrelid, ix.indkey::varchar,
-                  ix.indoption::varchar, i.reloptions, am.amname
+                  ix.indoption::varchar, i.reloptions, am.amname,
+                  %s as indnkeyatts
               FROM
                   pg_class t
                         join pg_index ix on t.oid = ix.indrelid
@@ -3319,7 +3321,11 @@ class PGDialect(default.DefaultDialect):
               ORDER BY
                   t.relname,
                   i.relname
-            """
+            """ % (
+                "ix.indnkeyatts"
+                if self.server_version_info >= (11, 0)
+                else "NULL",
+            )
 
         t = sql.text(IDX_SQL).columns(
             relname=sqltypes.Unicode, attname=sqltypes.Unicode
@@ -3342,6 +3348,7 @@ class PGDialect(default.DefaultDialect):
                 idx_option,
                 options,
                 amname,
+                indnkeyatts,
             ) = row
 
             if expr:
@@ -3365,7 +3372,10 @@ class PGDialect(default.DefaultDialect):
             if col is not None:
                 index["cols"][col_num] = col
             if not has_idx:
-                index["key"] = [int(k.strip()) for k in idx_key.split()]
+                idx_keys = idx_key.split()
+                if indnkeyatts is not None:
+                    idx_keys = idx_keys[:indnkeyatts]
+                index["key"] = [int(k.strip()) for k in idx_keys]
 
                 # (new in pg 8.3)
                 # "pg_index.indoption" is list of ints, one per column/expr.
diff --git a/test/dialect/postgresql/test_reflection.py b/test/dialect/postgresql/test_reflection.py
index 374bda5658..9b0ffdaa61 100644
--- a/test/dialect/postgresql/test_reflection.py
+++ b/test/dialect/postgresql/test_reflection.py
@@ -1110,6 +1110,30 @@ class ReflectionTest(fixtures.TestBase):
                 "gin",
             )
 
+    @testing.fails_if("postgresql < 11.0", "indnkeyatts not supported")
+    @testing.provide_metadata
+    def test_index_reflection_with_include(self):
+        """reflect indexes with include set"""
+
+        metadata = self.metadata
+
+        Table(
+            "t",
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("x", ARRAY(Integer)),
+            Column("name", String(20)),
+        )
+        metadata.create_all()
+        with testing.db.connect().execution_options(autocommit=True) as conn:
+            conn.execute("CREATE INDEX idx1 ON t (x) INCLUDE (name)")
+
+            ind = testing.db.dialect.get_indexes(conn, "t", None)
+            eq_(
+                ind,
+                [{"unique": False, "column_names": ["x"], "name": "idx1"}],
+            )
+
     @testing.provide_metadata
     def test_foreign_key_option_inspection(self):
         metadata = self.metadata
-- 
2.47.3