11# safe_ddl.py – helpers for idempotent Alembic operations
22# Place this inside migrations/utils/ or another importable location
33
4- from typing import Any , Optional
4+ from typing import Any , List , Optional
55
66import sqlalchemy as sa
77from alembic import op
1616 "safe_alter_column" ,
1717 "rename_column_if_exists" ,
1818 "safe_execute_sql" ,
19+ # FK helpers (new)
20+ "has_fk" ,
21+ "drop_fk_if_exists" ,
22+ "create_fk_if_not_exists" ,
23+ "replace_fk" ,
1924]
2025
2126# ─────────────────────────────────────────────────────────────────────────────
@@ -42,6 +47,115 @@ def has_column(table_name: str, column_name: str) -> bool:
4247# Some scripts import 'column_exists'; keep it as an alias for compatibility.
4348column_exists = has_column
4449
50+
51+ # ─────────────────────────────────────────────────────────────────────────────
52+ # FK inspection helpers (new)
53+ # ─────────────────────────────────────────────────────────────────────────────
54+
55+
56+ def has_fk (table_name : str , constraint_name : str ) -> bool :
57+ """
58+ Return True if a foreign-key constraint with the given name exists on the table.
59+
60+ Uses information_schema so constraint names are authoritative (not inferred
61+ from SQLAlchemy metadata, which may differ from what MySQL actually stored).
62+ """
63+ if not has_table (table_name ):
64+ return False
65+ bind = op .get_bind ()
66+ insp = inspect (bind )
67+ fks = insp .get_foreign_keys (table_name )
68+ return any (fk .get ("name" ) == constraint_name for fk in fks )
69+
70+
71+ def drop_fk_if_exists (table_name : str , constraint_name : str ) -> None :
72+ """
73+ Drop a foreign-key constraint by name only if it currently exists.
74+ Safe to call on a table that has already been migrated.
75+ """
76+ if not has_table (table_name ):
77+ _log (f"⚠️ Skipped drop FK – table not found: { table_name } " )
78+ return
79+
80+ if has_fk (table_name , constraint_name ):
81+ op .drop_constraint (constraint_name , table_name , type_ = "foreignkey" )
82+ _log (f"🗑️ Dropped FK: { table_name } .{ constraint_name } " )
83+ else :
84+ _log (f"⚠️ Skipped drop FK – constraint not found: { table_name } .{ constraint_name } " )
85+
86+
87+ def create_fk_if_not_exists (
88+ constraint_name : str ,
89+ source_table : str ,
90+ referent_table : str ,
91+ local_cols : List [str ],
92+ remote_cols : List [str ],
93+ * ,
94+ ondelete : Optional [str ] = None ,
95+ onupdate : Optional [str ] = None ,
96+ ) -> None :
97+ """
98+ Create a foreign-key constraint only if it does not already exist.
99+ Safe to call on a table that has already been migrated.
100+ """
101+ if not has_table (source_table ):
102+ _log (f"⚠️ Skipped create FK – source table not found: { source_table } " )
103+ return
104+
105+ if has_fk (source_table , constraint_name ):
106+ _log (f"⚠️ Skipped create FK – constraint already exists: { source_table } .{ constraint_name } " )
107+ return
108+
109+ op .create_foreign_key (
110+ constraint_name ,
111+ source_table ,
112+ referent_table ,
113+ local_cols ,
114+ remote_cols ,
115+ ondelete = ondelete ,
116+ onupdate = onupdate ,
117+ )
118+ rule = f" ON DELETE { ondelete } " if ondelete else ""
119+ _log (f"✅ Created FK: { source_table } .{ constraint_name } → { referent_table } { rule } " )
120+
121+
122+ def replace_fk (
123+ constraint_name : str ,
124+ source_table : str ,
125+ referent_table : str ,
126+ local_cols : List [str ],
127+ remote_cols : List [str ],
128+ * ,
129+ ondelete : Optional [str ] = None ,
130+ onupdate : Optional [str ] = None ,
131+ ) -> None :
132+ """
133+ Atomically drop-and-recreate a FK constraint with new delete/update rules.
134+
135+ This is the canonical helper for FK migrations:
136+ - Drops only if the constraint currently exists (idempotent on upgrade).
137+ - Creates only if it does not yet exist (idempotent on re-run).
138+
139+ Usage:
140+ replace_fk(
141+ "files_ibfk_1",
142+ "files", "users",
143+ ["user_id"], ["id"],
144+ ondelete="CASCADE",
145+ )
146+ """
147+ drop_fk_if_exists (source_table , constraint_name )
148+ create_fk_if_not_exists (
149+ constraint_name ,
150+ source_table ,
151+ referent_table ,
152+ local_cols ,
153+ remote_cols ,
154+ ondelete = ondelete ,
155+ onupdate = onupdate ,
156+ )
157+
158+
45159# ─────────────────────────────────────────────────────────────────────────────
46160# Column-level operations
47161# ─────────────────────────────────────────────────────────────────────────────
@@ -50,46 +164,47 @@ def has_column(table_name: str, column_name: str) -> bool:
50164def add_column_if_missing (table_name : str , column : sa .Column ) -> None :
51165 """Add a column only if the table exists and the column is missing."""
52166 if not has_table (table_name ):
53- _log (f"⚠️ Skipped add column – table not found: { table_name } " )
167+ _log (f"⚠️ Skipped add column – table not found: { table_name } " )
54168 return
55169
56170 if not has_column (table_name , column .name ):
57171 with op .batch_alter_table (table_name ) as batch_op :
58172 batch_op .add_column (column )
59- _log (f"✅ Added column: { table_name } .{ column .name } " )
173+ _log (f"✅ Added column: { table_name } .{ column .name } " )
60174 else :
61- _log (f"⚠️ Skipped add column – column already exists: { table_name } .{ column .name } " )
175+ _log (f"⚠️ Skipped add column – column already exists: { table_name } .{ column .name } " )
62176
63177
64178def drop_column_if_exists (table_name : str , column_name : str ) -> None :
65179 """Drop a column only if both the table and column exist."""
66180 if not has_table (table_name ):
67- _log (f"⚠️ Skipped drop column – table not found: { table_name } " )
181+ _log (f"⚠️ Skipped drop column – table not found: { table_name } " )
68182 return
69183
70184 if has_column (table_name , column_name ):
71185 with op .batch_alter_table (table_name ) as batch_op :
72186 batch_op .drop_column (column_name )
73- _log (f"🗑️ Dropped column: { table_name } .{ column_name } " )
187+ _log (f"🗑️ Dropped column: { table_name } .{ column_name } " )
74188 else :
75- _log (f"⚠️ Skipped drop column – column already absent: { table_name } .{ column_name } " )
189+ _log (f"⚠️ Skipped drop column – column already absent: { table_name } .{ column_name } " )
76190
77191
78192def safe_alter_column (table_name : str , column_name : str , ** kwargs : Any ) -> None :
79193 """
80194 Alter a column only if both the table and column exist.
81- kwargs are passed directly to batch_op.alter_column (e.g., nullable=..., type_=..., existing_type=...).
195+ kwargs are passed directly to batch_op.alter_column
196+ (e.g., nullable=..., type_=..., existing_type=...).
82197 """
83198 if not has_table (table_name ):
84- _log (f"⚠️ Skipped alter column – table not found: { table_name } " )
199+ _log (f"⚠️ Skipped alter column – table not found: { table_name } " )
85200 return
86201
87202 if has_column (table_name , column_name ):
88203 with op .batch_alter_table (table_name ) as batch_op :
89204 batch_op .alter_column (column_name , ** kwargs )
90- _log (f"✏️ Altered column: { table_name } .{ column_name } " )
205+ _log (f"✏️ Altered column: { table_name } .{ column_name } " )
91206 else :
92- _log (f"⚠️ Skipped alter column – column not found: { table_name } .{ column_name } " )
207+ _log (f"⚠️ Skipped alter column – column not found: { table_name } .{ column_name } " )
93208
94209
95210def rename_column_if_exists (
@@ -101,19 +216,19 @@ def rename_column_if_exists(
101216 existing_nullable : Optional [bool ] = None ,
102217) -> None :
103218 """
104- Rename a column (old_name -> new_name) only if table exists, old exists, and new doesn't.
105- MySQL/MariaDB require existing_type/existing_nullable for ALTER in some cases; pass them if known .
219+ Rename a column (old_name -> new_name) only if table exists,
220+ old exists, and new doesn't .
106221 """
107222 if not has_table (table_name ):
108- _log (f"⚠️ Skipped rename – table not found: { table_name } " )
223+ _log (f"⚠️ Skipped rename – table not found: { table_name } " )
109224 return
110225
111226 if not has_column (table_name , old_name ):
112- _log (f"⚠️ Skipped rename – source column not found: { table_name } .{ old_name } " )
227+ _log (f"⚠️ Skipped rename – source column not found: { table_name } .{ old_name } " )
113228 return
114229
115230 if has_column (table_name , new_name ):
116- _log (f"⚠️ Skipped rename – target already exists: { table_name } .{ new_name } " )
231+ _log (f"⚠️ Skipped rename – target already exists: { table_name } .{ new_name } " )
117232 return
118233
119234 with op .batch_alter_table (table_name ) as batch_op :
@@ -123,7 +238,7 @@ def rename_column_if_exists(
123238 existing_type = existing_type ,
124239 existing_nullable = existing_nullable ,
125240 )
126- _log (f"🔤 Renamed column: { table_name } .{ old_name } → { new_name } " )
241+ _log (f"🔤 Renamed column: { table_name } .{ old_name } → { new_name } " )
127242
128243
129244# ─────────────────────────────────────────────────────────────────────────────
@@ -133,11 +248,12 @@ def rename_column_if_exists(
133248
134249def safe_execute_sql (sql : str ) -> None :
135250 """
136- Execute raw SQL and print a short log. Use for guarded UPDATE backfills, etc.
251+ Execute raw SQL and print a short log.
252+ Use for guarded UPDATE backfills etc.
137253 Caller is responsible for table/column checks.
138254 """
139255 bind = op .get_bind ()
140- _log (f"🛠️ Executing SQL: { sql } " )
256+ _log (f"🛠️ Executing SQL: { sql } " )
141257 bind .exec_driver_sql (sql )
142258
143259
0 commit comments