Copilot commented on code in PR #64838:
URL: https://github.com/apache/airflow/pull/64838#discussion_r3045122164
##########
airflow-core/src/airflow/migrations/utils.py:
##########
@@ -54,9 +54,10 @@ def get_mssql_table_constraints(conn, table_name) ->
dict[str, dict[str, list[st
@contextmanager
def disable_sqlite_fkeys(op):
if op.get_bind().dialect.name == "sqlite":
- op.execute("PRAGMA foreign_keys=off")
- yield op
- op.execute("PRAGMA foreign_keys=on")
+ with contextlib.ExitStack() as exit_stack:
+ op.execute("PRAGMA foreign_keys=off")
+ exit_stack.callback(op.execute, "PRAGMA foreign_keys=on")
+ yield op
Review Comment:
This change makes `disable_sqlite_fkeys()` exception-safe for migrations
that already use the helper, but there are still migration scripts that toggle
`PRAGMA foreign_keys` manually without a `try/finally`/ExitStack (e.g.
`0082_3_1_0_make_bundle_name_not_nullable.py` upgrade path). To fully deliver
on the goal of preventing connections from being left with foreign keys
disabled after a failed migration, consider updating the remaining manual
PRAGMA toggles to use this helper (or otherwise ensure they restore
`foreign_keys=ON` on error).
##########
airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py:
##########
@@ -188,19 +189,30 @@ def upgrade() -> None:
dialect_name = conn.dialect.name
if dialect_name == "sqlite":
- conn.execute(sa.text("PRAGMA foreign_keys=OFF"))
-
- with op.batch_alter_table("deadline", schema=None) as batch_op:
- batch_op.add_column(sa.Column("deadline_alert_id", sa.Uuid(),
nullable=True))
- batch_op.add_column(sa.Column("created_at", UtcDateTime,
nullable=True))
- batch_op.add_column(sa.Column("last_updated_at", UtcDateTime,
nullable=True))
- batch_op.create_foreign_key(
- batch_op.f("deadline_deadline_alert_id_fkey"),
- "deadline_alert",
- ["deadline_alert_id"],
- ["id"],
- ondelete="SET NULL",
- )
+ with disable_sqlite_fkeys(op):
Review Comment:
In the SQLite path, `disable_sqlite_fkeys()` currently only wraps the first
`batch_alter_table("deadline", ...)` block, but the subsequent
`op.batch_alter_table("deadline", ...)` (alter_column calls) and
`op.batch_alter_table("deadline_alert", ...)` blocks later in `upgrade()` run
with foreign keys re-enabled. Previously, foreign keys were disabled for the
whole sequence of SQLite batch operations; re-enabling early can break SQLite
batch migrations (and defeats the purpose of reliably keeping FK checks
disabled during all batch operations). Consider wrapping all SQLite batch
operations in a single `with disable_sqlite_fkeys(op):` block (and removing the
duplicated `else` branch by using the helper unconditionally as a noop for
non-SQLite).
--
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]
