ashb commented on code in PR #54508:
URL: https://github.com/apache/airflow/pull/54508#discussion_r2276711426
##########
airflow-core/src/airflow/migrations/versions/0082_3_1_0_make_bundle_name_not_nullable.py:
##########
@@ -99,11 +99,32 @@ def upgrade():
def downgrade():
"""Make bundle_name nullable."""
- with op.batch_alter_table("dag", schema=None) as batch_op:
- batch_op.drop_constraint(batch_op.f("dag_bundle_name_fkey"),
type_="foreignkey")
+ dialect_name = op.get_bind().dialect.name
- batch_op.alter_column("bundle_name", nullable=True,
existing_type=StringID())
- with op.batch_alter_table("dag", schema=None) as batch_op:
- batch_op.create_foreign_key(
- batch_op.f("dag_bundle_name_fkey"), "dag_bundle", ["bundle_name"],
["name"]
- )
+ if dialect_name == "sqlite":
+ # SQLite requires foreign key constraints to be disabled during batch
operations
+ conn = op.get_bind()
+ conn.execute(text("PRAGMA foreign_keys=OFF"))
+
+ try:
+ with op.batch_alter_table("dag", schema=None) as batch_op:
+ batch_op.drop_constraint(batch_op.f("dag_bundle_name_fkey"),
type_="foreignkey")
+ batch_op.alter_column("bundle_name", nullable=True,
existing_type=StringID())
+
+ with op.batch_alter_table("dag", schema=None) as batch_op:
+ batch_op.create_foreign_key(
+ batch_op.f("dag_bundle_name_fkey"), "dag_bundle",
["bundle_name"], ["name"]
+ )
+ finally:
+ # Always re-enable foreign key constraints
+ conn.execute(text("PRAGMA foreign_keys=ON"))
+ else:
+ # For PostgreSQL and MySQL, use the original logic
+ with op.batch_alter_table("dag", schema=None) as batch_op:
+ batch_op.drop_constraint(batch_op.f("dag_bundle_name_fkey"),
type_="foreignkey")
+ batch_op.alter_column("bundle_name", nullable=True,
existing_type=StringID())
+
+ with op.batch_alter_table("dag", schema=None) as batch_op:
+ batch_op.create_foreign_key(
+ batch_op.f("dag_bundle_name_fkey"), "dag_bundle",
["bundle_name"], ["name"]
+ )
Review Comment:
Am I reading this correctly, that the only difference between sqlite and
other branches is wrapping it in a pragma to disable the FKs? If so A better
pattern might be:
```suggestion
exitstack = contextlib.ExitStack
if dialect_name == "sqlite":
# SQLite requires foreign key constraints to be disabled during
batch operations
conn = op.get_bind()
conn.execute(text("PRAGMA foreign_keys=OFF"))
exitstack.callback(conn.execute, text("PRAGMA foreign_keys=ON"))
with exitstack:
with op.batch_alter_table("dag", schema=None) as batch_op:
batch_op.drop_constraint(batch_op.f("dag_bundle_name_fkey"),
type_="foreignkey")
batch_op.alter_column("bundle_name", nullable=True,
existing_type=StringID())
with op.batch_alter_table("dag", schema=None) as batch_op:
batch_op.create_foreign_key(
batch_op.f("dag_bundle_name_fkey"), "dag_bundle",
["bundle_name"], ["name"]
)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
