diff --git a/alembic/ddl/postgresql.py b/alembic/ddl/postgresql.py index cc03f453..54f195dd 100644 --- a/alembic/ddl/postgresql.py +++ b/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 ( diff --git a/alembic/op.pyi b/alembic/op.pyi index 8b347d21..6708e983 100644 --- a/alembic/op.pyi +++ b/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 ``_``. 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 ``_``. 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. """ diff --git a/alembic/operations/base.py b/alembic/operations/base.py index 67aefdd4..dc7a76a9 100644 --- a/alembic/operations/base.py +++ b/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 ``_``. 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 ``_``. 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 ... diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py index c474c07b..8b2a9e34 100644 --- a/alembic/operations/ops.py +++ b/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 ``_``. 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 ``_``. 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. """ diff --git a/docs/build/unreleased/1723.rst b/docs/build/unreleased/1723.rst new file mode 100644 index 00000000..b1cc2bb3 --- /dev/null +++ b/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. diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index d0304651..d9b2c810 100644 --- a/tests/test_postgresql.py +++ b/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)