Fix postgresql_include in create_index

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.

Copied from e01041b and 174cb4b.

Fixes: #874
Closes: #875
Pull-request: https://github.com/sqlalchemy/alembic/pull/875
Pull-request-sha: 2f9b51890a1296146bcea842dc8bbee53d1d11f5

Change-Id: I68cc2a75dafd3297c5121c5125ff9c463c74f777

diff --git a/alembic/ddl/postgresql.py b/alembic/ddl/postgresql.py
index c894649a5917dde7877f5b9bb9bc26975ca9c2a1..9fb9ac980582c1b35c25e8f8af9f8e12adab0aac 100644 (file)
--- 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 (file)
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 500678a6b545c318d5e5402d3e7163cf21d1b71a..b41c383be6046f0354ca0fe508b28bb8516333cf 100644 (file)
--- 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)