This is an automated email from the ASF dual-hosted git repository.
kaxil 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 568d6b3cce8 Fix airflow db clean silently skipping three tables when
scoped to Dags (#73173)
568d6b3cce8 is described below
commit 568d6b3cce81544226ea9979a20dd3abcde5640e
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Sep 16 18:11:46 2026 +0100
Fix airflow db clean silently skipping three tables when scoped to Dags
(#73173)
---
airflow-core/src/airflow/utils/db_cleanup.py | 101 ++++++++++++++++-
airflow-core/tests/unit/utils/test_db_cleanup.py | 136 ++++++++++++++++++++++-
2 files changed, 226 insertions(+), 11 deletions(-)
diff --git a/airflow-core/src/airflow/utils/db_cleanup.py
b/airflow-core/src/airflow/utils/db_cleanup.py
index e956fcf2e52..53e1971b8e7 100644
--- a/airflow-core/src/airflow/utils/db_cleanup.py
+++ b/airflow-core/src/airflow/utils/db_cleanup.py
@@ -82,6 +82,28 @@ def _format_table_name(schema: str | None, table: str) ->
str:
return table
[email protected](frozen=True)
+class _IndirectDagScope:
+ """
+ Describe how to scope a table to a Dag when it carries no ``dag_id`` of
its own.
+
+ ``--dag-ids`` / ``--exclude-dag-ids`` normally filter on a column of the
table being
+ cleaned. A table that reaches its Dag only through a foreign key needs the
filter
+ expressed as a subquery against the referenced table instead.
+
+ :param fk_column: the foreign key on the table being cleaned; must be
listed in
+ ``extra_columns`` so it is present on the constructed table
+ :param referenced_table: the table ``fk_column`` points at
+ :param referenced_pk_column: the primary key of ``referenced_table`` that
``fk_column`` matches
+ :param referenced_dag_id_column: the Dag id column on ``referenced_table``
+ """
+
+ fk_column: str
+ referenced_table: str
+ referenced_pk_column: str = "id"
+ referenced_dag_id_column: str = "dag_id"
+
+
@dataclasses.dataclass
class _TableConfig:
"""
@@ -97,6 +119,9 @@ class _TableConfig:
:param keep_last_group_by: if keeping the last record, can keep the last
record for each group
:param dependent_tables: list of tables which have FK relationship with
this table
:param extra_filters: SQLAlchemy expressions ANDed with the recency
filter; referenced columns must be in ``extra_columns``.
+ :param dag_id_scope: how to apply ``--dag-ids`` / ``--exclude-dag-ids`` to
a table that has no
+ Dag id column of its own and reaches its Dag through a foreign key.
Mutually exclusive with
+ ``dag_id_column_name``.
:param skip_if_referenced: list of ``(referencing_table, fk_column)``
pairs whose FK points at this
table's ``referenced_pk_column``. A row that is still referenced by
any of these is excluded from
deletion. This avoids issuing deletes that would violate an ``ON
DELETE RESTRICT`` foreign key
@@ -109,6 +134,7 @@ class _TableConfig:
recency_column_name: str
extra_columns: list[str] | None = None
dag_id_column_name: str | None = None
+ dag_id_scope: _IndirectDagScope | None = None
keep_last: bool = False
keep_last_filters: Any | None = None
keep_last_group_by: Any | None = None
@@ -145,6 +171,20 @@ class _TableConfig:
schema=self.schema_name,
)
+ if self.dag_id_scope is not None:
+ if self.dag_id_column_name is not None:
+ raise ValueError(
+ f"_TableConfig for table {self.table_name!r} sets both
dag_id_column_name and "
+ f"dag_id_scope; a table is scoped to a Dag either by its
own column or through a "
+ f"foreign key, not both."
+ )
+ if self.dag_id_scope.fk_column not in self.orm_model.c.keys():
+ raise ValueError(
+ f"_TableConfig for table {self.table_name!r} sets
dag_id_scope but its "
+ f"fk_column {self.dag_id_scope.fk_column!r} is not one of
its columns; "
+ f"add {self.dag_id_scope.fk_column!r} to extra_columns."
+ )
+
# skip_if_referenced filters on referenced_pk_column, which must be a
column of orm_model
# (added via extra_columns). Fail fast with a clear message instead of
a cryptic KeyError
# raised later when _build_query evaluates
base_table.c[referenced_pk_column].
@@ -163,7 +203,12 @@ class _TableConfig:
return {
"table": self.table_name,
"recency_column": str(self.recency_column),
- "dag_id_column": str(self.dag_id_column),
+ "dag_id_column": (
+ f"{self.dag_id_scope.fk_column} -> "
+
f"{self.dag_id_scope.referenced_table}.{self.dag_id_scope.referenced_dag_id_column}"
+ if self.dag_id_scope is not None
+ else str(self.dag_id_column)
+ ),
"keep_last": self.keep_last,
"keep_last_filters": [str(x) for x in self.keep_last_filters] if
self.keep_last_filters else None,
"keep_last_group_by": str(self.keep_last_group_by),
@@ -188,7 +233,11 @@ config_list: list[_TableConfig] = [
keep_last_group_by=["dag_id"],
dependent_tables=["task_instance", "task_state_store", "deadline"],
),
- _TableConfig(table_name="asset_event", recency_column_name="timestamp",
dag_id_column_name="dag_id"),
+ # asset_event has never had a dag_id; the producing Dag is source_dag_id,
and it is NULL for
+ # events that no task produced (an API-created event, or a watcher).
+ _TableConfig(
+ table_name="asset_event", recency_column_name="timestamp",
dag_id_column_name="source_dag_id"
+ ),
# Carries no foreign key, so rows are left behind when the partition Dag
run they describe
# is cascade-deleted with its dag_run. Only such orphans may be purged:
rows whose partition
# Dag run still exists are the evidence the scheduler evaluates to decide
when that pending
@@ -219,7 +268,14 @@ config_list: list[_TableConfig] = [
recency_column_name="expires_at",
dag_id_column_name="dag_id",
),
- _TableConfig(table_name="task_reschedule",
recency_column_name="start_date", dag_id_column_name="dag_id"),
+ # task_reschedule.dag_id was dropped in 3.0.0; a reschedule now reaches
its Dag through its
+ # task instance. ti_id is NOT NULL, so no row is unattributed.
+ _TableConfig(
+ table_name="task_reschedule",
+ recency_column_name="start_date",
+ extra_columns=["ti_id"],
+ dag_id_scope=_IndirectDagScope(fk_column="ti_id",
referenced_table="task_instance"),
+ ),
_TableConfig(table_name="xcom", recency_column_name="timestamp",
dag_id_column_name="dag_id"),
_TableConfig(table_name="_xcom_archive", recency_column_name="timestamp",
dag_id_column_name="dag_id"),
_TableConfig(
@@ -268,7 +324,16 @@ config_list: list[_TableConfig] = [
# and are cleaned. dag_run.created_dag_version_id is ON DELETE SET
NULL, so it does not block.
skip_if_referenced=[("task_instance", "dag_version_id")],
),
- _TableConfig(table_name="deadline", recency_column_name="deadline_time",
dag_id_column_name="dag_id"),
+ # deadline.dag_id was dropped in 3.1.0; a deadline now reaches its Dag
through its dag run.
+ # The scope has to follow, because this table is cleaned as a dependent of
dag_run precisely so
+ # its rows are archived before the ON DELETE CASCADE removes them --
leaving it unscoped would
+ # purge deadlines for Dags whose runs --dag-ids / --exclude-dag-ids is
preserving.
+ _TableConfig(
+ table_name="deadline",
+ recency_column_name="deadline_time",
+ extra_columns=["dagrun_id"],
+ dag_id_scope=_IndirectDagScope(fk_column="dagrun_id",
referenced_table="dag_run"),
+ ),
_TableConfig(table_name="revoked_token", recency_column_name="exp"),
_TableConfig(
table_name="connection_test_request",
@@ -479,6 +544,7 @@ def _build_query(
clean_before_timestamp: DateTime,
session: Session,
dag_id_column=None,
+ dag_id_scope: _IndirectDagScope | None = None,
dag_ids: list[str] | None = None,
exclude_dag_ids: list[str] | None = None,
extra_filters: list[Any] | None = None,
@@ -517,7 +583,30 @@ def _build_query(
if dag_ids:
conditions.append(base_table_dag_id_col.in_(dag_ids))
if exclude_dag_ids:
- conditions.append(base_table_dag_id_col.not_in(exclude_dag_ids))
+ # A NULL dag id belongs to no Dag, so it is not one of the
excluded Dags' rows and stays
+ # eligible. NOT IN alone would yield NULL for it and silently
retain it forever -- which
+ # is every `job` row, since core never sets Job.dag_id.
+ conditions.append(
+ or_(base_table_dag_id_col.is_(None),
base_table_dag_id_col.not_in(exclude_dag_ids))
+ )
+ elif (dag_ids or exclude_dag_ids) and dag_id_scope is not None:
+ fk_col = base_table.c[dag_id_scope.fk_column]
+ referenced = table(
+ dag_id_scope.referenced_table,
+ column(dag_id_scope.referenced_pk_column),
+ column(dag_id_scope.referenced_dag_id_column),
+ )
+
+ def _rows_for(target_dag_ids: list[str]):
+ return
select(referenced.c[dag_id_scope.referenced_pk_column]).where(
+
referenced.c[dag_id_scope.referenced_dag_id_column].in_(target_dag_ids)
+ )
+
+ if dag_ids:
+ conditions.append(fk_col.in_(_rows_for(dag_ids)))
+ if exclude_dag_ids:
+ # NULL-safe for the same reason as the direct-column branch above.
+ conditions.append(or_(fk_col.is_(None),
fk_col.not_in(_rows_for(exclude_dag_ids))))
if keep_last:
max_date_col_name = "max_date_per_group"
@@ -549,6 +638,7 @@ def _cleanup_table(
keep_last_group_by,
clean_before_timestamp: DateTime,
dag_id_column=None,
+ dag_id_scope: _IndirectDagScope | None = None,
dag_ids=None,
exclude_dag_ids=None,
dry_run: bool = True,
@@ -568,6 +658,7 @@ def _cleanup_table(
orm_model=orm_model,
recency_column=recency_column,
dag_id_column=dag_id_column,
+ dag_id_scope=dag_id_scope,
dag_ids=dag_ids,
exclude_dag_ids=exclude_dag_ids,
keep_last=keep_last,
diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py
b/airflow-core/tests/unit/utils/test_db_cleanup.py
index f67b45d748c..01e3ade8ce3 100644
--- a/airflow-core/tests/unit/utils/test_db_cleanup.py
+++ b/airflow-core/tests/unit/utils/test_db_cleanup.py
@@ -38,9 +38,11 @@ from airflow.exceptions import AirflowException
from airflow.models import DagModel, DagRun, TaskInstance
from airflow.models.dag_version import DagVersion
from airflow.models.dagbundle import DagBundleModel
+from airflow.models.deadline import Deadline
from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.task_state_store import TaskStateStoreModel
from airflow.providers.standard.operators.python import PythonOperator
+from airflow.sdk.definitions.callback import AsyncCallback
from airflow.serialization.serialized_objects import LazyDeserializedDAG
from airflow.utils.db_cleanup import (
ARCHIVE_TABLE_PREFIX,
@@ -52,6 +54,7 @@ from airflow.utils.db_cleanup import (
_dump_table_to_file,
_effective_table_names,
_get_archived_table_names,
+ _IndirectDagScope,
_TableConfig,
config_dict,
drop_archived_tables,
@@ -63,8 +66,10 @@ from airflow.utils.types import DagRunType
from tests_common.test_utils.db import (
clear_db_assets,
+ clear_db_callbacks,
clear_db_dag_bundles,
clear_db_dags,
+ clear_db_deadline,
clear_db_runs,
drop_tables_with_prefix,
)
@@ -447,6 +452,88 @@ class TestDBCleanup:
f"Expected {expected_remaining_dag_ids} to remain, but got
{remaining_dag_ids}"
)
+ @pytest.mark.parametrize(
+ ("dag_ids", "exclude_dag_ids", "expected_remaining"),
+ [
+ pytest.param(["dag1"], None, {"dag2", None},
id="include_scopes_through_dag_run"),
+ pytest.param(None, ["dag1"], {"dag1"},
id="exclude_keeps_only_that_dags_deadlines"),
+ pytest.param(["dag1", "dag2"], ["dag2"], {"dag2", None},
id="include_and_exclude"),
+ pytest.param(None, None, set(), id="unfiltered_purges_everything"),
+ ],
+ )
+ def test_deadline_cleanup_is_scoped_through_its_dag_run(
+ self, dag_ids, exclude_dag_ids, expected_remaining
+ ):
+ """
+ ``deadline`` has carried no ``dag_id`` since 3.1.0, so it is scoped
via ``dagrun_id``.
+
+ A deadline with no dag run belongs to no Dag: ``--dag-ids`` must not
claim it, and
+ ``--exclude-dag-ids`` must not shield it.
+ """
+ base_date = pendulum.DateTime(2022, 1, 1,
tzinfo=pendulum.timezone("UTC"))
+
+ with create_session() as session:
+ bundle_name = "testing"
+ session.add(DagBundleModel(name=bundle_name))
+ session.flush()
+
+ runs_by_dag = {}
+ for dag_id in ["dag1", "dag2"]:
+ dag = DAG(dag_id=dag_id)
+ session.add(DagModel(dag_id=dag_id, bundle_name=bundle_name))
+
SerializedDagModel.write_dag(LazyDeserializedDAG.from_dag(dag),
bundle_name=bundle_name)
+ dag_run = DagRun(
+ dag_id,
+ run_id=f"{dag_id}_run",
+ run_type=DagRunType.MANUAL,
+ start_date=base_date,
+ )
+ session.add(dag_run)
+ session.flush()
+ runs_by_dag[dag_id] = dag_run.id
+
+ for run_id in runs_by_dag.values():
+ session.add(
+ Deadline(
+ deadline_time=base_date,
+
callback=AsyncCallback("tests.unit.models.test_deadline.callback_for_deadline"),
+ dagrun_id=run_id,
+ deadline_alert_id=None,
+ )
+ )
+ # A deadline attached to no dag run at all.
+ session.add(
+ Deadline(
+ deadline_time=base_date,
+
callback=AsyncCallback("tests.unit.models.test_deadline.callback_for_deadline"),
+ dagrun_id=None,
+ deadline_alert_id=None,
+ )
+ )
+ session.commit()
+
+ run_cleanup(
+ clean_before_timestamp=base_date.add(days=10),
+ table_names=["deadline"],
+ dag_ids=dag_ids,
+ exclude_dag_ids=exclude_dag_ids,
+ dry_run=False,
+ confirm=False,
+ session=session,
+ )
+
+ run_id_to_dag = {run_id: dag_id for dag_id, run_id in
runs_by_dag.items()}
+ remaining = {
+ run_id_to_dag.get(deadline.dagrun_id) for deadline in
session.scalars(select(Deadline)).all()
+ }
+ assert remaining == expected_remaining
+
+ # The deadline with no dag run has nothing to cascade from, and
callback rows are only
+ # reachable deadline -> callback, so neither is removed by the dag/run
clears in
+ # clean_database. Left behind they leak into whatever runs next
against this database.
+ clear_db_deadline()
+ clear_db_callbacks()
+
@pytest.mark.parametrize(
("skip_archive", "expected_archives"),
[pytest.param(True, 0, id="skip_archive"), pytest.param(False, 1,
id="do_archive")],
@@ -562,6 +649,28 @@ class TestDBCleanup:
# "id" intentionally omitted from extra_columns
)
+ def test_table_config_rejects_both_dag_id_column_and_scope(self):
+ """A table reaches its Dag one way or the other, so naming both ways
must fail fast."""
+ with pytest.raises(ValueError, match="both dag_id_column_name and"):
+ _TableConfig(
+ table_name="deadline",
+ recency_column_name="deadline_time",
+ dag_id_column_name="dag_id",
+ dag_id_scope=_IndirectDagScope(fk_column="dagrun_id",
referenced_table="dag_run"),
+ # present, so this fails on the conflict rather than on a
missing fk_column
+ extra_columns=["dagrun_id"],
+ )
+
+ def test_table_config_dag_id_scope_requires_fk_column(self):
+ """A dag_id_scope whose fk_column is not selected must fail fast at
construction."""
+ with pytest.raises(ValueError, match="fk_column"):
+ _TableConfig(
+ table_name="deadline",
+ recency_column_name="deadline_time",
+ dag_id_scope=_IndirectDagScope(fk_column="dagrun_id",
referenced_table="dag_run"),
+ # "dagrun_id" intentionally omitted from extra_columns
+ )
+
def test_do_delete_rolls_back_before_drop_on_failure(self):
session = MagicMock(spec=Session)
session.get_bind.return_value.dialect.name = "mysql"
@@ -897,13 +1006,32 @@ class TestDBCleanup:
assert set(all_models) - exclusion_list.union(config_dict) == set()
assert exclusion_list.isdisjoint(config_dict)
- def test_no_failure_warnings(self):
+ @pytest.mark.parametrize(
+ ("dag_ids", "exclude_dag_ids"),
+ [
+ pytest.param(None, None, id="unfiltered"),
+ pytest.param(["some_dag"], None, id="include"),
+ pytest.param(None, ["some_dag"], id="exclude"),
+ pytest.param(["some_dag"], ["other_dag"],
id="include_and_exclude"),
+ ],
+ )
+ def test_no_failure_warnings(self, dag_ids, exclude_dag_ids):
"""
Ensure every table we have configured (and that is present in the db)
can be cleaned successfully.
For example, this checks that the recency column is actually a column.
+
+ The Dag-scoped parametrizations matter as much as the unfiltered one:
the Dag filter is the
+ only thing that dereferences ``dag_id_column_name`` /
``dag_id_scope``, so a config naming a
+ column a migration has dropped compiles fine without them. Three
tables drifted that way
+ across 3.0.0 and 3.1.0 before this was covered.
"""
with patch("airflow.utils.db_cleanup.logger") as mock_logger:
- run_cleanup(clean_before_timestamp=timezone.utcnow(), dry_run=True)
+ run_cleanup(
+ clean_before_timestamp=timezone.utcnow(),
+ dag_ids=dag_ids,
+ exclude_dag_ids=exclude_dag_ids,
+ dry_run=True,
+ )
for call in mock_logger.warning.call_args_list:
assert "Encountered error when attempting to clean table" not
in str(call)
@@ -1474,8 +1602,6 @@ class TestCallbackCleanup:
assert bool(survived) is should_survive
def test_unfired_deadline_callback_and_its_deadline_survive(self,
dag_maker):
- from airflow.models.deadline import Deadline
- from airflow.sdk.definitions.callback import AsyncCallback
from airflow.utils.state import CallbackState
old = pendulum.now(tz="UTC").subtract(days=30)
@@ -1556,8 +1682,6 @@ class TestCallbackCleanup:
cutoff = pendulum.now(tz="UTC").subtract(days=1)
future = pendulum.now(tz="UTC").add(days=365)
- from airflow.models.deadline import Deadline
-
with create_session() as session:
callback_id = self._add_callback(session, "scheduled", old)
session.execute(
