From 55d2300b68e8ef4bcc268526228f25f6ee957cc7 Mon Sep 17 00:00:00 2001
From: Justin Malin <justin@joincandidhealth.com>
Date: Fri, 6 Jun 2025 20:36:53 +0000
Subject: [PATCH] [Fixes #1671] Passthrough `dialect_kwargs` to
 `ops.create_foreign_key()` to enable `postgresql_not_valid`

---
 alembic/autogenerate/render.py | 33 ++++++++++++++++++++-------------
 tests/test_postgresql.py       |  5 +++++
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 4bad437a..d5ca2f4f 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -18,6 +18,7 @@ from mako.pygen import PythonPrinter
 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
@@ -121,6 +122,15 @@ def _render_cmd_body(
     return buf.getvalue()
 
 
+def _render_dialect_kwargs_items(
+        autogen_context: AutogenContext, dialect_kwargs: _DialectArgView
+) -> list[str]:
+    return [
+        f"{key}={_render_potential_expr(val, autogen_context)}"
+        for key, val in dialect_kwargs.items()
+    ]
+
+
 def render_op(
     autogen_context: AutogenContext, op: ops.MigrateOperation
 ) -> List[str]:
@@ -303,15 +313,6 @@ def _drop_table(autogen_context: AutogenContext, op: ops.DropTableOp) -> str:
     return text
 
 
-def _render_dialect_kwargs_items(
-    autogen_context: AutogenContext, item: DialectKWArgs
-) -> list[str]:
-    return [
-        f"{key}={_render_potential_expr(val, autogen_context)}"
-        for key, val in item.dialect_kwargs.items()
-    ]
-
-
 @renderers.dispatch_for(ops.CreateIndexOp)
 def _add_index(autogen_context: AutogenContext, op: ops.CreateIndexOp) -> str:
     index = op.to_index()
@@ -331,7 +332,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 +366,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 % {
@@ -418,9 +419,15 @@ def _add_fk_constraint(
             if value is not None:
                 args.append("%s=%r" % (k, value))
 
-    return "%(prefix)screate_foreign_key(%(args)s)" % {
+
+    # op.to_constraint() may fail with `multiple values for argument 'name'`.
+    dialect_kwarg_view = getattr(op.kw, "dialect_kwargs", {})
+    dialect_kwargs = _render_dialect_kwargs_items(autogen_context, dialect_kwarg_view)
+
+    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 +671,7 @@ 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))]
diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py
index 44f422dd..289090df 100644
--- a/tests/test_postgresql.py
+++ b/tests/test_postgresql.py
@@ -133,6 +133,11 @@ class PostgresqlOpTest(TestBase):
         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")
-- 
2.47.3