Fix SQLite pragma teardown in DB migrations by iting0321 · Pull Request #64838 · apache/airflow

iting0321 · 2026-04-07T12:50:09Z

What

This change improves migration utility reliability across SQLite, MySQL, and offline Alembic flows. It adds a shared SQLite foreign-key helper, updates affected migrations to use it, and hardens related helpers for offline-safe dialect detection and retry-safe MySQL stored procedures.

Why

To prevent SQLite migrations from leaving foreign keys disabled after failures, avoid offline Alembic crashes during dialect detection, and make MySQL migration helper reruns safer after interrupted executions.

Was generative AI tooling used to co-author this PR?

Yes (please specify the tool below) Claude Code following the guidelines

boring-cyborg · 2026-04-07T12:50:13Z

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:

Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it.
Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
Be sure to read the Airflow Coding style.
Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack

Copilot

Pull request overview

This PR aims to make SQLite foreign-key PRAGMA handling in Alembic migrations reliable by ensuring foreign_keys=ON is always restored (even when a migration errors), and by updating selected migration scripts to use a shared helper along with new regression tests.

Changes:

Make disable_sqlite_fkeys() exception-safe so PRAGMA teardown happens even on errors.
Update several existing migration scripts to use disable_sqlite_fkeys() instead of manual PRAGMA toggles.
Add unit tests covering success/exception/no-op behavior of disable_sqlite_fkeys().

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
airflow-core/src/airflow/migrations/utils.py	Makes the SQLite FK-disabling helper exception-safe via `ExitStack` callback teardown.
airflow-core/tests/unit/migrations/test_migration_utils.py	Adds regression tests ensuring PRAGMA is restored on both success and exception.
airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py	Switches SQLite FK handling to the shared helper (but currently narrows the disable scope).
airflow-core/src/airflow/migrations/versions/0084_3_1_0_add_last_parse_duration_to_dag_model.py	Replaces manual PRAGMA toggling with the helper in downgrade.
airflow-core/src/airflow/migrations/versions/0041_3_0_0_rename_dataset_as_asset.py	Replaces manual PRAGMA toggling with the helper in `_drop_fkey_if_exists()`.
airflow-core/src/airflow/migrations/versions/0017_2_9_2_fix_inconsistency_between_ORM_and_migration_files.py	Replaces manual PRAGMA toggling with the helper for SQLite table-recreation paths.

Copilot

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Copilot

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Copilot

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.

jason810496

Thanks for the improvement.

Copilot

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

potiuk · 2026-04-22T19:31:20Z

@iting0321 Converting to draft — this PR doesn't yet meet our Pull Request quality criteria.

❌ Unresolved review comments (5 threads from maintainers): please walk through each unresolved review thread. Even if a suggestion looks incorrect or irrelevant — and some of them will be, especially any comments left by automated reviewers like GitHub Copilot — it is still the author's responsibility to respond: apply the fix, reply in-thread with a brief explanation of why the suggestion does not apply, or resolve the thread if the feedback is no longer relevant. Leaving threads unaddressed for weeks blocks the PR from moving forward.

See the linked criteria for how to fix each item, then mark the PR "Ready for review". This is not a rejection — just an invitation to bring the PR up to standard. No rush.

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

potiuk · 2026-04-22T22:19:10Z

Quick follow-up to the triage comment above — one clarification on the "Unresolved review comments" item:

Once you believe a thread has been addressed — whether by pushing a fix, or by replying in-thread with an explanation of why the suggestion doesn't apply — please mark the thread as resolved yourself by clicking the "Resolve conversation" button at the bottom of each thread. Reviewers don't auto-close their own threads, so an addressed-but-unresolved thread reads as "still waiting on the author" and keeps the PR from moving forward. The author doing the resolve-click is the expected convention on this project.

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

… for utility functions

… enhancements

…del migration

… tests

…greSQL

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 7 comments.

…MySQL in migration utilities

…r for table creation and data insertion

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

+        conn = op.get_bind()
+
+        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 op.batch_alter_table("deadline", schema=None) as batch_op:
-        batch_op.alter_column("created_at", existing_type=UtcDateTime, nullable=False)
-        batch_op.alter_column("last_updated_at", existing_type=UtcDateTime, nullable=False)
-
-    with op.batch_alter_table("deadline_alert", schema=None) as batch_op:
-        batch_op.create_foreign_key(
-            batch_op.f("deadline_alert_serialized_dag_id_fkey"),
-            "serialized_dag",
-            ["serialized_dag_id"],
-            ["id"],
-            ondelete="CASCADE",
+        # For migration/backcompat purposes if no timestamp is there from the migration, use now()
+        # then lock the columns down so all new entries require the timestamps to be provided.
+        now = timezone.utcnow()
+        conn.execute(
+            sa.text("""
+                UPDATE deadline
+                SET created_at = :now, last_updated_at = :now
+                WHERE created_at IS NULL OR last_updated_at IS NULL
+            """),
+            {"now": now},
        )


+    if dialect_name == "postgresql":
+        cols_array = ", ".join(f"'{c}'" for c in columns)
+        op.execute(


 def _drop_fkey_if_exists(table, constraint_name):
    conn = op.get_bind()
    dialect_name = conn.dialect.name



Lee-W · 2026-05-11T09:35:38Z

+        self.executed.append(statement)
+
+
+def test_disable_sqlite_fkeys_restores_pragma_on_success() -> None:


We no longer need this. we have a unit test for it in airflow-core/tests/unit/migrations/test_migration_utils.py

we can merge tests there if needed

Lee-W · 2026-05-11T09:39:46Z

    return contextlib.nullcontext()
+
+
+def create_index_if_not_exists(op, index_name, table_name, columns, unique=False) -> None:


is this still used somewhere?

Lee-W · 2026-05-11T09:41:56Z

+    else:
+        # SQLite requires foreign key constraints to be disabled during batch operations.
+        with disable_sqlite_fkeys(op):
+            try:


should we use ignore_sqlite_value_error?

Lee-W · 2026-05-11T09:47:42Z

+        op.create_index(index_name, table_name, columns, unique=unique, if_not_exists=True)
+
+
+def drop_index_if_exists(op, index_name, table_name) -> None:


is this still used somewhere?

Lee-W · 2026-05-11T09:47:46Z

+        op.drop_index(index_name, table_name=table_name, if_exists=True)
+
+
+def drop_unique_constraints_on_columns(op, table_name, columns) -> None:


is this still used somewhere?

Lee-W · 2026-05-11T09:47:54Z

+                    batch_op.drop_constraint(uq["name"], type_="unique")
+
+
+def drop_unique_constraint_if_exists(op, table_name, constraint_name) -> None:


is this still used somewhere?

iting0321 requested a review from ephraimbuddy as a code owner April 7, 2026 12:50

boring-cyborg Bot added area:db-migrations PRs with DB migration area:deadline-alerts AIP-86 (former AIP-57) labels Apr 7, 2026

jason810496 requested review from Copilot and jason810496 April 7, 2026 12:52

Copilot started reviewing on behalf of jason810496 April 7, 2026 12:53 View session

Copilot AI reviewed Apr 7, 2026

View reviewed changes

Comment thread airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py Outdated

Comment thread airflow-core/src/airflow/migrations/utils.py

jason810496 requested a review from Copilot April 7, 2026 13:45

Copilot started reviewing on behalf of jason810496 April 7, 2026 13:46 View session

Copilot AI reviewed Apr 7, 2026

View reviewed changes

Comment thread airflow-core/src/airflow/migrations/versions/0082_3_1_0_make_bundle_name_not_nullable.py Outdated

Comment thread airflow-core/src/airflow/migrations/versions/0084_3_1_0_add_last_parse_duration_to_dag_model.py Outdated

jason810496 requested review from Copilot and removed request for ephraimbuddy April 7, 2026 13:58

Copilot started reviewing on behalf of jason810496 April 7, 2026 14:00 View session

Copilot AI reviewed Apr 7, 2026

View reviewed changes

Comment thread airflow-core/tests/unit/migrations/test_migration_utils.py Outdated

jason810496 requested a review from Copilot April 7, 2026 14:24

Copilot started reviewing on behalf of jason810496 April 7, 2026 14:25 View session

Copilot AI reviewed Apr 7, 2026

View reviewed changes

jason810496 approved these changes Apr 8, 2026

View reviewed changes

jason810496 requested review from Lee-W and ephraimbuddy April 8, 2026 13:10

Lee-W reviewed Apr 10, 2026

View reviewed changes

kaxil requested a review from Copilot April 10, 2026 19:55

Copilot AI reviewed Apr 10, 2026

View reviewed changes

potiuk marked this pull request as draft April 22, 2026 19:31

iting0321 added 3 commits April 29, 2026 21:30

Refactor SQLite foreign key handling in migrations and add unit tests…

e4975bf

… for utility functions

Refactor SQLite handling in migration scripts and improve deadline UI…

df68dcc

… enhancements

Refactor SQLite foreign key handling in downgrade function for DAG mo…

069734e

…del migration

iting0321 added 4 commits April 29, 2026 21:30

Fix disable_sqlite_fkeys to use get_dialect_name and add offline mode…

c4cdfcf

… tests

Refactor foreign key handling in migration scripts for MySQL and Post…

6273f3a

…greSQL

fix get_dialect_name function

0831db4

restore 0017_2_9_2_fix_inconsistency_between_ORM_and_migration_files.py

ddc8bfe

iting0321 force-pushed the db-migration branch from bbca803 to ddc8bfe Compare April 29, 2026 13:32

iting0321 added 2 commits April 30, 2026 22:42

Moved the disable_sqlite_fkeys unit tests into a new file

9cfa896

Add missing import for contextlib in migration utilities

4d81786

jason810496 requested a review from Copilot May 1, 2026 14:54

Copilot started reviewing on behalf of jason810496 May 1, 2026 14:55 View session

Copilot AI reviewed May 1, 2026

View reviewed changes

iting0321 added 2 commits May 4, 2026 14:20

Refactor SQLite error handling and add DROP PROCEDURE statements for …

a832aca

…MySQL in migration utilities

Refactor upgrade functions to use disable_sqlite_fkeys context manage…

17f7557

…r for table creation and data insertion

iting0321 marked this pull request as ready for review May 4, 2026 14:46

potiuk added the ready for maintainer review Set after triaging when all criteria pass. label May 5, 2026

jason810496 requested a review from Copilot May 9, 2026 11:09

Copilot started reviewing on behalf of jason810496 May 9, 2026 11:10 View session

Copilot AI reviewed May 9, 2026

View reviewed changes

Lee-W reviewed May 11, 2026

View reviewed changes

		self.executed.append(statement)


		def test_disable_sqlite_fkeys_restores_pragma_on_success() -> None:

		return contextlib.nullcontext()


		def create_index_if_not_exists(op, index_name, table_name, columns, unique=False) -> None:

		op.create_index(index_name, table_name, columns, unique=unique, if_not_exists=True)


		def drop_index_if_exists(op, index_name, table_name) -> None:

		op.drop_index(index_name, table_name=table_name, if_exists=True)


		def drop_unique_constraints_on_columns(op, table_name, columns) -> None:

		batch_op.drop_constraint(uq["name"], type_="unique")


		def drop_unique_constraint_if_exists(op, table_name, constraint_name) -> None:

Conversation

iting0321 commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Why

Uh oh!

boring-cyborg Bot commented Apr 7, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

jason810496 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

potiuk commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

potiuk commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Lee-W May 11, 2026

Choose a reason for hiding this comment

Uh oh!

Lee-W May 11, 2026

Choose a reason for hiding this comment

Uh oh!

Lee-W May 11, 2026

Choose a reason for hiding this comment

Uh oh!

Lee-W May 11, 2026

Choose a reason for hiding this comment

Uh oh!

Lee-W May 11, 2026

Choose a reason for hiding this comment

Uh oh!

Lee-W May 11, 2026

Choose a reason for hiding this comment

Uh oh!

iting0321 commented Apr 7, 2026 •

edited

Loading

potiuk commented Apr 22, 2026 •

edited

Loading

potiuk commented Apr 22, 2026 •

edited

Loading