Merge "[Fixes #1671] Passthrough dialect_kwargs to ops.create_foreign_key()" into main

Author: zzzeek

alembic/autogenerate/render.py

@@ -18,6 +18,7 @@
 from sqlalchemy import schema as sa_schema
 from sqlalchemy import sql
 from sqlalchemy import types as sqltypes
+from sqlalchemy.sql.base import _DialectArgView
 from sqlalchemy.sql.elements import conv
 from sqlalchemy.sql.elements import Label
 from sqlalchemy.sql.elements import quoted_name
@@ -31,7 +32,6 @@
 
     from sqlalchemy import Computed
     from sqlalchemy import Identity
-    from sqlalchemy.sql.base import DialectKWArgs
     from sqlalchemy.sql.elements import ColumnElement
     from sqlalchemy.sql.elements import TextClause
     from sqlalchemy.sql.schema import CheckConstraint
@@ -304,11 +304,11 @@ def _drop_table(autogen_context: AutogenContext, op: ops.DropTableOp) -> str:
 
 
 def _render_dialect_kwargs_items(
-    autogen_context: AutogenContext, item: DialectKWArgs
+    autogen_context: AutogenContext, dialect_kwargs: _DialectArgView
 ) -> list[str]:
     return [
         f"{key}={_render_potential_expr(val, autogen_context)}"
-        for key, val in item.dialect_kwargs.items()
+        for key, val in dialect_kwargs.items()
     ]
 
 
@@ -331,7 +331,7 @@ def _add_index(autogen_context: AutogenContext, op: ops.CreateIndexOp) -> str:
 
     assert index.table is not None
 
-    opts = _render_dialect_kwargs_items(autogen_context, index)
+    opts = _render_dialect_kwargs_items(autogen_context, index.dialect_kwargs)
     if op.if_not_exists is not None:
         opts.append("if_not_exists=%r" % bool(op.if_not_exists))
     text = tmpl % {
@@ -365,7 +365,7 @@ def _drop_index(autogen_context: AutogenContext, op: ops.DropIndexOp) -> str:
             "%(prefix)sdrop_index(%(name)r, "
             "table_name=%(table_name)r%(schema)s%(kwargs)s)"
         )
-    opts = _render_dialect_kwargs_items(autogen_context, index)
+    opts = _render_dialect_kwargs_items(autogen_context, index.dialect_kwargs)
     if op.if_exists is not None:
         opts.append("if_exists=%r" % bool(op.if_exists))
     text = tmpl % {
@@ -389,6 +389,7 @@ def _add_unique_constraint(
 def _add_fk_constraint(
     autogen_context: AutogenContext, op: ops.CreateForeignKeyOp
 ) -> str:
+    constraint = op.to_constraint()
     args = [repr(_render_gen_name(autogen_context, op.constraint_name))]
     if not autogen_context._has_batch:
         args.append(repr(_ident(op.source_table)))
@@ -418,9 +419,16 @@ def _add_fk_constraint(
             if value is not None:
                 args.append("%s=%r" % (k, value))
 
-    return "%(prefix)screate_foreign_key(%(args)s)" % {
+    dialect_kwargs = _render_dialect_kwargs_items(
+        autogen_context, constraint.dialect_kwargs
+    )
+
+    return "%(prefix)screate_foreign_key(%(args)s%(dialect_kwargs)s)" % {
         "prefix": _alembic_autogenerate_prefix(autogen_context),
         "args": ", ".join(args),
+        "dialect_kwargs": (
+            ", " + ", ".join(dialect_kwargs) if dialect_kwargs else ""
+        ),
     }
 
 
@@ -664,7 +672,9 @@ def _uq_constraint(
         opts.append(
             ("name", _render_gen_name(autogen_context, constraint.name))
         )
-    dialect_options = _render_dialect_kwargs_items(autogen_context, constraint)
+    dialect_options = _render_dialect_kwargs_items(
+        autogen_context, constraint.dialect_kwargs
+    )
 
     if alter:
         args = [repr(_render_gen_name(autogen_context, constraint.name))]

docs/build/unreleased/1672.rst

@@ -0,0 +1,10 @@
+.. change::
+    :tags: usecase, operations
+    :tickets: 1671
+    Fixed an issue where dialect-specific keyword arguments, dialect_kwargs, were
+    not passed through when using the op.create_foreign_key() operation. This
+    prevented the use of backend-specific foreign key options, such as
+    postgresql_not_valid for PostgreSQL constraints. The renderer for
+    ops.CreateForeignKeyOp now correctly includes these arguments, aligning its
+    behavior with other constraint operations.
+    Pull request courtesy of Justin Malin.

tests/test_autogen_render.py

@@ -314,6 +314,24 @@ def test_render_drop_index_custom_kwarg(self):
             "somedialect_foobar='option')",
         )
 
+    def test_add_fk_constraint__dialect_kwargs(self):
+        t1 = self.table()
+        t2 = self.table()
+        item = ForeignKeyConstraint(
+            [t1.c.id], [t2.c.id], name="fk", postgresql_not_valid=True
+        )
+        fk_obj = ops.CreateForeignKeyOp.from_constraint(item)
+
+        eq_ignore_whitespace(
+            re.sub(
+                r"u'",
+                "'",
+                autogenerate.render_op_text(self.autogen_context, fk_obj),
+            ),
+            "op.create_foreign_key('fk', 'test', 'test', ['id'], ['id'], "
+            "postgresql_not_valid=True)",
+        )
+
     def test_drop_index_batch(self):
         """
         autogenerate.render._drop_index

tests/test_postgresql.py

@@ -133,6 +133,16 @@ def test_create_index_if_not_exists(self):
         op.create_index("i", "t", ["c1", "c2"], if_not_exists=True)
         context.assert_("CREATE INDEX IF NOT EXISTS i ON t (c1, c2)")
 
+    def test_create_fk_postgresql_not_valid(self):
+        context = op_fixture("postgresql")
+        op.create_foreign_key(
+            "i", "t1", "t2", ["c1"], ["c2"], postgresql_not_valid=True
+        )
+        context.assert_(
+            "ALTER TABLE t1 ADD CONSTRAINT i FOREIGN KEY(c1) "
+            "REFERENCES t2 (c2) NOT VALID"
+        )
+
     @config.combinations("include_table", "no_table", argnames="include_table")
     def test_drop_index_postgresql_concurrently(self, include_table):
         context = op_fixture("postgresql")