This is an automated email from the ASF dual-hosted git repository.
ferruzzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b68ea7890b9 Fix airflow db clean IntegrityError when a dag_version is
referenced by a task_instance (#66350)
b68ea7890b9 is described below
commit b68ea7890b918844f68d34d506b4a442ef3042ee
Author: Jakub Matyszewski <[email protected]>
AuthorDate: Tue Sep 22 21:17:30 2026 +0300
Fix airflow db clean IntegrityError when a dag_version is referenced by a
task_instance (#66350)
* Fix IntegrityError in airflow db clean when dag_version is referenced by
task_instance
The ON DELETE RESTRICT FK from task_instance.dag_version_id means that a
dag_version row cannot be deleted while any task instance still references
it.
The skip_if_referenced filter in _build_query correctly excluded such rows
from
the archive (SELECT) step, but _do_delete issued the DELETE using only a PK
join against the archive table — with no re-check of the FK constraint.
In production Airflow (running while db clean executes) a new task_instance
can
be created referencing a dag_version between the archive INSERT commit and
the
DELETE, causing a psycopg2 ForeignKeyViolation (IntegrityError) and aborting
the cleanup command.
Re-apply the same NOT EXISTS guard on the DELETE WHERE clause so rows that
have
become referenced since the archive was created are silently skipped
instead of
failing the FK.
Also add:
- test_do_delete_skip_if_referenced_guards_against_race: directly exercises
the
_do_delete FK guard by inserting a TI after the archive is created.
- test_dag_id_column_name_matches_schema: regression guard that reflects
every
configured table and asserts dag_id_column_name is an actual DB column,
addressing reviewer feedback requesting protection against config/schema
drift.
* Remove redundant local import of _do_delete now exported at module level
* Fix airflow db clean failing when a dag filter is used
asset_event, task_reschedule and deadline were configured to filter on a
dag_id column that none of them has, so passing --dag-ids or
--exclude-dag-ids made the cleanup fail on exactly those tables while the
same command without a dag filter succeeded.
* Add newsfragment for db clean IntegrityError mitigation
* Fix skip_if_referenced race test to use _build_query
The regression test called _do_delete with a raw select(archive_table),
bypassing the NOT EXISTS guard that _build_query embeds in every SELECT.
That made continue on deleted == 0 loop forever: the unguarded SELECT
kept counting the archive row even though the DELETE kept refusing it.
Rewrite the test to build the query via _build_query, as _cleanup_table
does. The TI reference then keeps the SELECT count at zero from the start
so the loop exits without issuing any DELETE, which is the correct
invariant: referenced rows are excluded by the SELECT, not just the DELETE.
* Warn when FK guard skips rows and reproduce the race in the test
When skip_if_referenced causes the DELETE to skip rows, log a warning
so users can tell their dag_version rows were left behind due to live
task_instance references, rather than silently continuing.
The race regression test was not actually reproducing the race: it
inserted the TI before building the query, so the NOT EXISTS guard in
the SELECT returned zero rows immediately and _do_delete never reached
the DELETE at all. Rewrite the test using a patched reflect_tables to
inject the TI in the exact window between the archive CTAS commit and
the DELETE, which is the real race scenario. Add a ``raced`` sentinel
so the test fails fast if the injection point ever moves rather than
passing without testing anything. Also add a second DagVersion so
raced_old is actually eligible for deletion (dag_version is keep_last
per dag_id, so a lone version is always the survivor).
Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
---------
Co-authored-by: arose26 <[email protected]>
Co-authored-by: Claude Sonnet 4.6 <[email protected]>
---
airflow-core/newsfragments/66350.bugfix.rst | 1 +
airflow-core/src/airflow/utils/db_cleanup.py | 43 +++++++-
airflow-core/tests/unit/utils/test_db_cleanup.py | 127 ++++++++++++++++++++++-
3 files changed, 165 insertions(+), 6 deletions(-)
diff --git a/airflow-core/newsfragments/66350.bugfix.rst
b/airflow-core/newsfragments/66350.bugfix.rst
new file mode 100644
index 00000000000..b03dfd176e7
--- /dev/null
+++ b/airflow-core/newsfragments/66350.bugfix.rst
@@ -0,0 +1 @@
+Reduces the window where ``airflow db clean`` fails with an ``IntegrityError``
when a ``dag_version`` row was still referenced by a ``task_instance`` created
between the archive ``INSERT`` and the ``DELETE``.
diff --git a/airflow-core/src/airflow/utils/db_cleanup.py
b/airflow-core/src/airflow/utils/db_cleanup.py
index a6213d7de9e..9c56814632b 100644
--- a/airflow-core/src/airflow/utils/db_cleanup.py
+++ b/airflow-core/src/airflow/utils/db_cleanup.py
@@ -30,7 +30,7 @@ import os
from collections.abc import Generator
from contextlib import contextmanager, suppress
from types import SimpleNamespace
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
from sqlalchemy import and_, column, func, inspect, literal, literal_column,
or_, select, table, text
from sqlalchemy.exc import OperationalError, ProgrammingError
@@ -52,6 +52,7 @@ from airflow.utils.types import DagRunType
if TYPE_CHECKING:
from pendulum import DateTime
from sqlalchemy import Select
+ from sqlalchemy.engine import CursorResult
from sqlalchemy.orm import Session
from airflow.models import Base
@@ -383,7 +384,14 @@ def _dump_table_to_file(*, target_table: str, file_path:
str, export_format: str
def _do_delete(
- *, query: Select, orm_model: Base, skip_archive: bool, session: Session,
batch_size: int | None
+ *,
+ query: Select,
+ orm_model: Base,
+ skip_archive: bool,
+ session: Session,
+ batch_size: int | None,
+ skip_if_referenced: list[tuple[str, str]] | None = None,
+ referenced_pk_column: str = "id",
) -> None:
import itertools
import re
@@ -454,10 +462,37 @@ def _do_delete(
delete = source_table.delete().where(
and_(*[col == target_table.c[col.name] for col in
source_table.primary_key.columns])
)
+ # Re-apply skip_if_referenced on the DELETE to guard against a
race where a new
+ # referencing row is created after the archive INSERT committed
but before the DELETE
+ # runs. Without this the DELETE would violate the ON DELETE
RESTRICT FK and fail.
+ if skip_if_referenced:
+ pk_col = source_table.c[referenced_pk_column]
+ for referencing_table_name, fk_column in skip_if_referenced:
+ referencing = table(referencing_table_name,
column(fk_column))
+ delete = delete.where(
+ ~select(literal(1))
+ .select_from(referencing)
+ .where(referencing.c[fk_column] == pk_col)
+ .correlate(source_table)
+ .exists()
+ )
logger.debug("delete statement:\n%s", delete.compile())
- session.execute(delete)
+ deleted = cast("CursorResult", session.execute(delete)).rowcount
session.commit()
+ # A guarded DELETE (skip_if_referenced) may delete fewer rows than
the SELECT
+ # found. The SELECT includes the same NOT EXISTS guard, so the
skipped row is
+ # excluded on the next pass too and the loop drains naturally.
With --batch-size
+ # set, continuing lets subsequent batches clean rows unaffected by
the race.
+ if deleted == 0:
+ logger.warning(
+ "Some rows from %s are still referenced by another table
and were not "
+ "deleted; they remain in %s and will be retried on the
next cleanup run.",
+ source_table_name,
+ target_table_name if not skip_archive else "the archive
(which is being dropped)",
+ )
+ continue
+
except BaseException:
error_raised = True
# Roll back the failed transaction so its locks are released before
@@ -678,6 +713,8 @@ def _cleanup_table(
skip_archive=skip_archive,
session=session,
batch_size=batch_size,
+ skip_if_referenced=skip_if_referenced,
+ referenced_pk_column=referenced_pk_column,
)
session.commit()
diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py
b/airflow-core/tests/unit/utils/test_db_cleanup.py
index 01e3ade8ce3..da3e7bae4a2 100644
--- a/airflow-core/tests/unit/utils/test_db_cleanup.py
+++ b/airflow-core/tests/unit/utils/test_db_cleanup.py
@@ -27,7 +27,18 @@ from uuid import uuid4
import pendulum
import pytest
-from sqlalchemy import Column, Integer, MetaData, Table, func, insert,
inspect, literal, select, text
+from sqlalchemy import (
+ Column,
+ Integer,
+ MetaData,
+ Table,
+ func,
+ insert,
+ inspect,
+ literal,
+ select,
+ text,
+)
from sqlalchemy.exc import IntegrityError, OperationalError, SQLAlchemyError
from sqlalchemy.ext.declarative import DeclarativeMeta
from sqlalchemy.orm import Session
@@ -638,6 +649,99 @@ class TestDBCleanup:
assert latest_id in remaining # kept by keep_last
assert orphan_id not in remaining # old and unreferenced -> pruned
+ def test_do_delete_skip_if_referenced_guards_against_race(self):
+ """_do_delete must not issue a DELETE that violates an ON DELETE
RESTRICT FK.
+
+ Reproduces the real race: the dag_version row passes the SELECT filter
and is
+ archived, and only then does a task_instance referencing it appear.
The
+ skip_if_referenced guard on the DELETE must skip the row instead of
failing with
+ IntegrityError, and the loop must still drain because the next SELECT
pass
+ re-evaluates the same NOT EXISTS guard and excludes it.
+ """
+ from airflow.utils.db import reflect_tables
+
+ base_date = pendulum.DateTime(2020, 1, 1,
tzinfo=pendulum.timezone("UTC"))
+ bundle_name = f"race-test-{uuid4()}"
+ dag_id = f"race_dag_{uuid4()}"
+
+ with create_session() as session:
+ session.add(DagBundleModel(name=bundle_name))
+ session.flush()
+ session.add(DagModel(dag_id=dag_id, bundle_name=bundle_name))
+ session.flush()
+
+ raced_old = DagVersion(
+ dag_id=dag_id,
+ version_number=1,
+ bundle_name=bundle_name,
+ created_at=base_date,
+ last_updated=base_date,
+ )
+ # dag_version is keep_last per dag_id, so a lone version is always
the
+ # keep_last survivor and is never eligible for deletion. A
second, newer
+ # version takes that role and leaves raced_old as the deletion
candidate.
+ latest = DagVersion(
+ dag_id=dag_id,
+ version_number=2,
+ bundle_name=bundle_name,
+ created_at=base_date.add(minutes=1),
+ last_updated=base_date.add(minutes=1),
+ )
+ session.add_all([raced_old, latest])
+ session.flush()
+ raced_old_id, latest_id = raced_old.id, latest.id
+
+ # Query built while nothing references raced_old, so the first
SELECT pass
+ # returns it and _do_delete archives it.
+ cfg = config_dict["dag_version"]
+ query = _build_query(
+ **cfg.__dict__,
+ clean_before_timestamp=base_date.add(days=10),
+ session=session,
+ )
+
+ dag_run = DagRun(dag_id, run_id="race-run",
run_type=DagRunType.MANUAL, start_date=base_date)
+ ti = create_task_instance(
+ PythonOperator(task_id="dummy-task", python_callable=print),
+ run_id=dag_run.run_id,
+ dag_version_id=raced_old_id,
+ )
+ ti.dag_id = dag_id
+ ti.start_date = base_date
+
+ raced = False
+
+ def reflect_and_race(tables, session, **kwargs):
+ # _do_delete reflects both source and target right after
committing the
+ # archive CTAS and right before building the DELETE — this is
the only
+ # seam between the two that fits the race window. MySQL
reflects the
+ # target alone earlier in the same pass, so keying on the
two-table call
+ # covers both branches. If this call site moves, the test
stops
+ # reproducing the race and the ``raced`` assertion below will
catch it.
+ nonlocal raced
+ if not raced and len(tables) == 2:
+ raced = True
+ session.add_all([dag_run, ti])
+ session.commit()
+ return reflect_tables(tables, session, **kwargs)
+
+ with patch("airflow.utils.db_cleanup.reflect_tables",
side_effect=reflect_and_race):
+ _do_delete(
+ query=query,
+ orm_model=cfg.orm_model,
+ skip_archive=True,
+ session=session,
+ batch_size=None,
+ skip_if_referenced=cfg.skip_if_referenced,
+ referenced_pk_column=cfg.referenced_pk_column,
+ )
+
+ remaining =
set(session.scalars(select(DagVersion.id).where(DagVersion.dag_id ==
dag_id)).all())
+
+ assert raced, "the TI was never inserted mid-pass; the race was not
reproduced"
+ assert raced_old_id in remaining, "dag_version referenced by a
task_instance must not be deleted"
+ assert latest_id in remaining, "the keep_last survivor must not be
deleted"
+
def test_table_config_skip_if_referenced_requires_pk_column(self):
"""A misconfigured skip_if_referenced (pk not in columns) must fail
fast at construction."""
with pytest.raises(ValueError, match="referenced_pk_column"):
@@ -747,7 +851,7 @@ class TestDBCleanup:
session.get_bind.return_value.dialect.name = "mysql"
session.connection.return_value = object()
session.scalars.return_value.one.side_effect = [1, 0]
- session.execute.side_effect = [None, None, None]
+ session.execute.side_effect = [None, None, MagicMock(rowcount=1)]
metadata, source_table, target_table, query =
_build_do_delete_test_objects()
@@ -810,7 +914,7 @@ class TestDBCleanup:
session.get_bind.return_value.dialect.name = "mysql"
session.connection.return_value = object()
session.scalars.return_value.one.side_effect = [1, 0]
- session.execute.side_effect = [None, None, None]
+ session.execute.side_effect = [None, None, MagicMock(rowcount=1)]
metadata, source_table, target_table, query =
_build_do_delete_test_objects()
drop_failure = OperationalError("DROP TABLE", {}, Exception("disk
full"))
@@ -1006,6 +1110,23 @@ class TestDBCleanup:
assert set(all_models) - exclusion_list.union(config_dict) == set()
assert exclusion_list.isdisjoint(config_dict)
+ def test_dag_id_column_name_matches_schema(self):
+ """
+ Regression guard: every dag_id_column_name in config_dict must be an
actual column in its
+ database table, so that --dag-ids filtering never raises
UndefinedColumn.
+ """
+ with create_session() as session:
+ insp = inspect(session.bind)
+ existing_tables = set(insp.get_table_names())
+ for table_name, cfg in config_dict.items():
+ if cfg.dag_id_column_name is None or table_name not in
existing_tables:
+ continue
+ db_columns = {col["name"] for col in
insp.get_columns(table_name)}
+ assert cfg.dag_id_column_name in db_columns, (
+
f"config_dict[{table_name!r}].dag_id_column_name={cfg.dag_id_column_name!r} "
+ f"is not a column of table {table_name!r} in the database"
+ )
+
@pytest.mark.parametrize(
("dag_ids", "exclude_dag_ids"),
[
