Support postgresql_include for create_primary_key / create_unique_constraint - #1819
Conversation
…straint The postgresql_include option added in SQLAlchemy 2.0.41 was only handled for create_index. Passing it to create_primary_key raised a TypeError as the operation did not accept dialect-specific keyword arguments, and create_unique_constraint raised a KeyError during DDL compilation because the INCLUDE columns were absent from the synthetic table built for the ALTER statement. create_primary_key now accepts dialect-specific keyword arguments, and the PostgreSQL implementation appends any postgresql_include columns referenced by name to the constraint's table before emitting AddConstraint, mirroring the existing create_index handling.
MohammedAlkindi
left a comment
There was a problem hiding this comment.
Checked this out and ran it on Windows 11, CPython 3.13.13, against main at c116cbc0.
Test evidence:
pytest tests/test_postgresql.py |
result |
|---|---|
| this PR | 71 passed |
| product change reverted, new tests kept | 2 failed, 69 passed |
The two failures with the product reverted are exactly this PR's new tests, so they discriminate the change. No existing test moved. (These are op_fixture("postgresql") SQL-string assertions — no live server involved, which is what let me run them on a bare Windows box; the actual PostgreSQL DDL round-trip is CI's to confirm.)
The refactor is a genuine improvement over the pre-existing code, not just an extension. The old inline loop in create_index appended Column(col, NullType) without checking the element type; the extracted _ensure_include_columns adds isinstance(col, str) before the in table.c test, and guards table is None. Centralising it also means create_index and add_constraint can't drift apart, which is how this class of bug (#1723) happened in the first place.
Two small questions:
postgresql_includeon an index can containColumnobjects as well as names (SQLAlchemy accepts both). The newisinstance(col, str)skips non-string elements — correct for the append (a realColumnneeds no synthesizing), but worth confirming the INCLUDE clause still renders for aColumnelement passed tocreate_primary_key, since the test matrix here covers the string form only.const.dialect_kwargs.get("postgresql_include", None) or ()— theor ()folds an explicitly-empty list andNonetogether, which is fine today; if SQLAlchemy ever distinguishes "empty INCLUDE" from "no INCLUDE", this line will hide it. A comment or a plainifwould make the intent explicit.
The version-gate note (SQLAlchemy 2.0.41 for constraint-level postgresql_include) in the comment matches the changelog entry, and the docs update covers both new call sites.
Fixes #1723
Problem
postgresql_include(added in SQLAlchemy 2.0.41 forPrimaryKeyConstraint/UniqueConstraint) was only wired up forcreate_index. Using it via the op directives failed in two distinct ways:op.create_primary_key(..., postgresql_include=[...])raisedTypeError: create_primary_key() got an unexpected keyword argument 'postgresql_include'because the public operation method did not accept dialect-specific keyword arguments.op.create_unique_constraint(..., postgresql_include=[...])raisedKeyErrorduring DDL compilation: SQLAlchemy's_define_includelooks the INCLUDE columns up inconstraint.table.c, but those columns are not present on the synthetic table Alembic builds for theALTER ... ADD CONSTRAINTstatement.Fix
CreatePrimaryKeyOp.create_primary_keynow accepts**kw, so dialect-specific kwargs flow through to the constraint (the rest of the op pipeline —__init__,to_constraint,from_constraint, andSchemaObjects.primary_key_constraint— already passed them through).create_unique_constraintalready accepted**kw.PostgresqlImpl.add_constraintnow appends anypostgresql_includecolumns referenced by name to the constraint's synthetic table before emittingAddConstraint, mirroring the existing handling increate_index. The shared logic is extracted into a small_ensure_include_columnshelper used by bothcreate_indexandadd_constraint..pyi/base.pyoperation stubs viatools/write_pyi.py(validated bytests/test_stubs.py).docs/build/unreleased/1723.rst.Tests
Added
test_create_unique_constraint_postgresql_includeandtest_create_primary_key_postgresql_includetotests/test_postgresql.py, asserting the emitted DDL:Both fail on
main(TypeError / KeyError respectively) and pass with this change. Full suite: 1769 passed, 129 skipped (DB-specific). flake8, black, mypy, and the stub-consistency test all pass.