Merge "Fix postgresql_include in create_index"

3 files changed, 33 insertions, 0 deletions

diff --git a/alembic/ddl/postgresql.py b/alembic/ddl/postgresql.py
index c894649..9fb9ac9 100644
--- a/alembic/ddl/postgresql.py
+++ b/alembic/ddl/postgresql.py
@@ -16,6 +16,7 @@ from sqlalchemy import types as sqltypes
 from sqlalchemy.dialects.postgresql import BIGINT
 from sqlalchemy.dialects.postgresql import ExcludeConstraint
 from sqlalchemy.dialects.postgresql import INTEGER
+from sqlalchemy.schema import CreateIndex
 from sqlalchemy.sql.elements import ColumnClause
 from sqlalchemy.sql.elements import UnaryExpression
 from sqlalchemy.types import NULLTYPE
@@ -71,6 +72,16 @@ class PostgresqlImpl(DefaultImpl):
     )
     identity_attrs_ignore = ("on_null", "order")
 
+    def create_index(self, index):
+        # this likely defaults to None if not present, so get()
+        # should normally not return the default value.  being
+        # defensive in any case
+        postgresql_include = index.kwargs.get("postgresql_include", None) or ()
+        for col in postgresql_include:
+            if col not in index.table.c:
+                index.table.append_column(Column(col, sqltypes.NullType))
+        self._exec(CreateIndex(index))
+
     def prep_table_for_batch(self, batch_impl, table):
 
         for constraint in table.constraints:
diff --git a/docs/build/unreleased/874.rst b/docs/build/unreleased/874.rst
new file mode 100644
index 0000000..36a73ef
--- /dev/null
+++ b/docs/build/unreleased/874.rst
@@ -0,0 +1,9 @@
+.. change::
+   :tags: bug, postgresql
+   :tickets: 874
+
+   Fixed issue where usage of the PostgreSQL ``postgresql_include`` option
+   within a :meth:`.Operations.create_index` would raise a KeyError, as the
+   additional column(s) need to be added to the table object used by the
+   construct internally. The issue is equivalent to the SQL Server issue fixed
+   in :ticket:`513`. Pull request courtesy Steven Bronson.
diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py
index 500678a..b41c383 100644
--- a/tests/test_postgresql.py
+++ b/tests/test_postgresql.py
@@ -103,6 +103,19 @@ class PostgresqlOpTest(TestBase):
             "CREATE INDEX CONCURRENTLY geocoded ON locations (coordinates)"
         )
 
+    @config.requirements.sqlalchemy_14
+    def test_create_index_postgresql_include(self):
+        context = op_fixture("postgresql")
+        op.create_index(
+            "i", "t", ["c1", "c2"], unique=False, postgresql_include=["inc"]
+        )
+        context.assert_("CREATE INDEX i ON t (c1, c2) INCLUDE (inc)")
+
+    def test_create_index_postgresql_include_is_none(self):
+        context = op_fixture("postgresql")
+        op.create_index("i", "t", ["c1", "c2"], unique=False)
+        context.assert_("CREATE INDEX i ON t (c1, c2)")
+
     def test_drop_index_postgresql_concurrently(self):
         context = op_fixture("postgresql")
         op.drop_index("geocoded", "locations", postgresql_concurrently=True)