Support postgresql_include for create_primary_key / create_unique_constraint

alembic/ddl/postgresql.py

@@ -91,13 +91,37 @@ def create_index(self, index: Index, **kw: Any) -> None:
         # 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:  # type: ignore[union-attr]
-                index.table.append_column(  # type: ignore[union-attr]
-                    Column(col, sqltypes.NullType)
-                )
+        self._ensure_include_columns(index.table, postgresql_include)
         self._exec(CreateIndex(index, **kw))
 
+    def add_constraint(self, const: Any, **kw: Any) -> None:
+        # the PostgreSQL ``postgresql_include`` option, supported on
+        # PrimaryKeyConstraint / UniqueConstraint since SQLAlchemy 2.0.41,
+        # references columns by name that may not be present on the
+        # synthetic table alembic builds for an ALTER.  ensure they exist
+        # so that the INCLUDE clause can be rendered (#1723), mirroring the
+        # handling done in ``create_index``.
+        postgresql_include = (
+            const.dialect_kwargs.get("postgresql_include", None) or ()
+        )
+        self._ensure_include_columns(
+            getattr(const, "table", None), postgresql_include
+        )
+        super().add_constraint(const, **kw)
+
+    @staticmethod
+    def _ensure_include_columns(
+        table: Optional[Table], include_columns: Sequence[str]
+    ) -> None:
+        # append any ``postgresql_include`` columns referenced by name that
+        # are not already present on the synthetic table, so DDL compilation
+        # can resolve them when rendering the INCLUDE clause.
+        if table is None:
+            return
+        for col in include_columns:
+            if isinstance(col, str) and col not in table.c:
+                table.append_column(Column(col, sqltypes.NullType))
+
     def prep_table_for_batch(self, batch_impl, table):
         for constraint in table.constraints:
             if (

alembic/op.pyi

@@ -780,8 +780,9 @@ def create_primary_key(
     columns: List[str],
     *,
     schema: Optional[str] = None,
+    **kw: Any,
 ) -> None:
-    """Issue a "create primary key" instruction using the current
+    r"""Issue a "create primary key" instruction using the current
     migration context.
 
     e.g.::
@@ -813,6 +814,11 @@ def create_primary_key(
      quoting of the schema outside of the default behavior, use
      the SQLAlchemy construct
      :class:`~sqlalchemy.sql.elements.quoted_name`.
+    :param \**kw: Additional keyword arguments are dialect specific, and
+     passed in the form ``<dialectname>_<argname>``.  See the documentation
+     regarding an individual dialect at :ref:`dialect_toplevel` for
+     detail on documented arguments.  An example of a dialect-specific
+     option is the PostgreSQL ``postgresql_include`` argument.
 
     """
 
@@ -937,7 +943,7 @@ def create_unique_constraint(
     schema: Optional[str] = None,
     **kw: Any,
 ) -> Any:
-    """Issue a "create unique constraint" instruction using the
+    r"""Issue a "create unique constraint" instruction using the
     current migration context.
 
     e.g.::
@@ -972,6 +978,11 @@ def create_unique_constraint(
      quoting of the schema outside of the default behavior, use
      the SQLAlchemy construct
      :class:`~sqlalchemy.sql.elements.quoted_name`.
+    :param \**kw: Additional keyword arguments are dialect specific, and
+     passed in the form ``<dialectname>_<argname>``.  See the documentation
+     regarding an individual dialect at :ref:`dialect_toplevel` for
+     detail on documented arguments.  An example of a dialect-specific
+     option is the PostgreSQL ``postgresql_include`` argument.
 
     """
 

alembic/operations/base.py

@@ -1217,8 +1217,9 @@ def create_primary_key(
             columns: List[str],
             *,
             schema: Optional[str] = None,
+            **kw: Any,
         ) -> None:
-            """Issue a "create primary key" instruction using the current
+            r"""Issue a "create primary key" instruction using the current
             migration context.
 
             e.g.::
@@ -1250,6 +1251,11 @@ def create_primary_key(
              quoting of the schema outside of the default behavior, use
              the SQLAlchemy construct
              :class:`~sqlalchemy.sql.elements.quoted_name`.
+            :param \**kw: Additional keyword arguments are dialect specific, and
+             passed in the form ``<dialectname>_<argname>``.  See the documentation
+             regarding an individual dialect at :ref:`dialect_toplevel` for
+             detail on documented arguments.  An example of a dialect-specific
+             option is the PostgreSQL ``postgresql_include`` argument.
 
             """  # noqa: E501
             ...
@@ -1380,7 +1386,7 @@ def create_unique_constraint(
             schema: Optional[str] = None,
             **kw: Any,
         ) -> Any:
-            """Issue a "create unique constraint" instruction using the
+            r"""Issue a "create unique constraint" instruction using the
             current migration context.
 
             e.g.::
@@ -1415,6 +1421,11 @@ def create_unique_constraint(
              quoting of the schema outside of the default behavior, use
              the SQLAlchemy construct
              :class:`~sqlalchemy.sql.elements.quoted_name`.
+            :param \**kw: Additional keyword arguments are dialect specific, and
+             passed in the form ``<dialectname>_<argname>``.  See the documentation
+             regarding an individual dialect at :ref:`dialect_toplevel` for
+             detail on documented arguments.  An example of a dialect-specific
+             option is the PostgreSQL ``postgresql_include`` argument.
 
             """  # noqa: E501
             ...

alembic/operations/ops.py

@@ -320,6 +320,7 @@ def create_primary_key(
         columns: List[str],
         *,
         schema: Optional[str] = None,
+        **kw: Any,
     ) -> None:
         """Issue a "create primary key" instruction using the current
         migration context.
@@ -353,9 +354,14 @@ def create_primary_key(
          quoting of the schema outside of the default behavior, use
          the SQLAlchemy construct
          :class:`~sqlalchemy.sql.elements.quoted_name`.
+        :param \\**kw: Additional keyword arguments are dialect specific, and
+         passed in the form ``<dialectname>_<argname>``.  See the documentation
+         regarding an individual dialect at :ref:`dialect_toplevel` for
+         detail on documented arguments.  An example of a dialect-specific
+         option is the PostgreSQL ``postgresql_include`` argument.
 
         """
-        op = cls(constraint_name, table_name, columns, schema=schema)
+        op = cls(constraint_name, table_name, columns, schema=schema, **kw)
         return operations.invoke(op)
 
     @classmethod
@@ -490,6 +496,11 @@ def create_unique_constraint(
          quoting of the schema outside of the default behavior, use
          the SQLAlchemy construct
          :class:`~sqlalchemy.sql.elements.quoted_name`.
+        :param \\**kw: Additional keyword arguments are dialect specific, and
+         passed in the form ``<dialectname>_<argname>``.  See the documentation
+         regarding an individual dialect at :ref:`dialect_toplevel` for
+         detail on documented arguments.  An example of a dialect-specific
+         option is the PostgreSQL ``postgresql_include`` argument.
 
         """
 

docs/build/unreleased/1723.rst

@@ -0,0 +1,12 @@
+.. change::
+    :tags: usecase, postgresql
+    :tickets: 1723
+
+    Added support for the PostgreSQL ``postgresql_include`` option, added in
+    SQLAlchemy 2.0.41, to the :meth:`.Operations.create_primary_key` and
+    :meth:`.Operations.create_unique_constraint` operations.  The
+    ``postgresql_include`` columns referenced by name are now added to the
+    table that Alembic constructs for the ``ALTER`` statement so that the
+    ``INCLUDE`` clause renders correctly, mirroring the existing handling for
+    :meth:`.Operations.create_index`.  The :meth:`.Operations.create_primary_key`
+    method additionally now accepts dialect-specific keyword arguments.

tests/test_postgresql.py

@@ -135,6 +135,25 @@ def test_create_index_postgresql_include_is_none(self):
         op.create_index("i", "t", ["c1", "c2"], unique=False)
         context.assert_("CREATE INDEX i ON t (c1, c2)")
 
+    def test_create_unique_constraint_postgresql_include(self):
+        context = op_fixture("postgresql")
+        op.create_unique_constraint(
+            "uq_t", "t", ["c1", "c2"], postgresql_include=["inc"]
+        )
+        context.assert_(
+            "ALTER TABLE t ADD CONSTRAINT uq_t UNIQUE (c1, c2) INCLUDE (inc)"
+        )
+
+    def test_create_primary_key_postgresql_include(self):
+        context = op_fixture("postgresql")
+        op.create_primary_key(
+            "pk_t", "t", ["c1", "c2"], postgresql_include=["inc"]
+        )
+        context.assert_(
+            "ALTER TABLE t ADD CONSTRAINT pk_t PRIMARY KEY (c1, c2) "
+            "INCLUDE (inc)"
+        )
+
     def test_create_index_if_not_exists(self):
         context = op_fixture("postgresql")
         op.create_index("i", "t", ["c1", "c2"], if_not_exists=True)