Vamsi-klu commented on code in PR #66350:
URL: https://github.com/apache/airflow/pull/66350#discussion_r4002479035
##########
airflow-core/src/airflow/utils/db_cleanup.py:
##########
@@ -392,10 +417,31 @@ 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 = session.execute(delete).rowcount
session.commit()
+ # A guarded DELETE (skip_if_referenced) may affect fewer rows than
the SELECT
+ # saw. If it affects zero, the next SELECT would return the same
rows and we
+ # would archive them forever, so break out. This also caps the
loop in any
+ # edge case where the DELETE unexpectedly removes nothing.
+ if deleted == 0:
Review Comment:
**[P1] Do not terminate cleanup on one guarded miss**
Breaking here terminates the entire cleanup, not merely the current batch.
With `batch_size=1`, if a concurrent task instance starts referencing the
archived Dag version before the guarded DELETE, `rowcount` becomes zero and
every later eligible version is left untouched. The production `_build_query`
already re-evaluates the same `NOT EXISTS` guard, so another iteration would
exclude the newly referenced row and advance. Please re-query and cover this
with at least two candidates.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/src/airflow/utils/db_cleanup.py:
##########
@@ -519,6 +566,31 @@ def _build_query(
if exclude_dag_ids:
conditions.append(base_table_dag_id_col.not_in(exclude_dag_ids))
Review Comment:
**[P1] Make nullable Dag exclusion NULL-aware**
`AssetEvent.source_dag_id` is nullable, so `source_dag_id NOT IN (...)`
evaluates to UNKNOWN for source-less events. Those events belong to none of the
excluded Dags but are retained whenever `--exclude-dag-ids` is used. Make this
exclusion NULL-aware, consistently with the new `dag_id_via` path, and add a
populated regression case.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/src/airflow/utils/db_cleanup.py:
##########
@@ -392,10 +417,31 @@ 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 = session.execute(delete).rowcount
Review Comment:
**[P2] Archive contains rows that were not deleted**
The archive commits before the second reference check. A newly referenced
row can therefore be skipped by the DELETE while the default path retains its
copy in `_airflow_deleted__`. That makes archive exports inaccurate and lets a
later cleanup archive the same primary key again. Reconcile the archive with
the rows actually deleted.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/tests/unit/utils/test_db_cleanup.py:
##########
@@ -551,6 +668,83 @@ def
test_dag_version_cleanup_skips_versions_pinned_by_task_instance(self):
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.
+
+ Simulates a race where a dag_version row passes the SELECT filter (no
TI
+ references it at archive-creation time) but a TI referencing it is
inserted
+ before the DELETE runs. The skip_if_referenced guard on the DELETE
itself
+ must leave the row in place instead of failing with IntegrityError.
+ """
+ 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()
+
+ dv = DagVersion(
+ dag_id=dag_id,
+ version_number=1,
+ bundle_name=bundle_name,
+ created_at=base_date,
+ last_updated=base_date,
+ )
+ session.add(dv)
+ session.flush()
+ dv_id = dv.id
+
+ # Manually create an archive table containing this dag_version row,
+ # simulating the CTAS step that ran before the TI was inserted.
+ # Use SQLAlchemy's Uuid codec to insert so the id encoding matches
how
+ # DagVersion.id was stored (hyphenated native uuid on Postgres,
.hex on
+ # MySQL/SQLite); a raw f-string would use the hyphenated form
everywhere
+ # and match nothing on MySQL/SQLite.
+ archive_name = f"{ARCHIVE_TABLE_PREFIX}dag_version__race_test"
+ stmt = text(
+ f"CREATE TABLE {archive_name} AS SELECT * FROM dag_version
WHERE id = :dv_id"
+ ).bindparams(bindparam("dv_id", value=dv_id, type_=Uuid()))
+ session.execute(stmt)
+ session.commit()
+
+ # Now insert a TI referencing the dag_version (the "race").
+ 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=dv_id,
+ )
+ ti.dag_id = dag_id
+ ti.start_date = base_date
+ session.add_all([dag_run, ti])
+ session.commit()
+
+ # Build a select query that would return the row (simulating what
_build_query
+ # returned before the TI was inserted).
+ metadata = reflect_tables([archive_name, "dag_version"], session)
+ archive_table = metadata.tables[archive_name]
+ query = select(archive_table)
Review Comment:
**[P1] Synthetic query manufactures the alleged loop**
`select(archive_table)` never contains the production `skip_if_referenced`
predicate, so it keeps returning the same row after the reference appears and
manufactures the infinite loop that motivated the harmful break. Build this
query through `_build_query`, use two candidates with `batch_size=1`, inject
the reference after the first archive commit, and prove the referenced row
survives while the later candidate is deleted.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/tests/unit/utils/test_db_cleanup.py:
##########
@@ -447,6 +462,108 @@ def test_cleanup_with_dag_id_filtering(self, dag_ids,
exclude_dag_ids, expected_
f"Expected {expected_remaining_dag_ids} to remain, but got
{remaining_dag_ids}"
)
+ @pytest.mark.parametrize(
+ ("dag_ids", "exclude_dag_ids"),
+ [
+ pytest.param(["dag1"], None, id="include"),
+ pytest.param(None, ["dag1"], id="exclude"),
+ ],
+ )
+ def test_cleanup_dag_filtering_on_tables_without_their_own_dag_id(self,
dag_ids, exclude_dag_ids):
+ """asset_event, task_reschedule and deadline have no dag_id column of
their own."""
+ with create_session() as session:
+ run_cleanup(
+ clean_before_timestamp=pendulum.DateTime(2022, 1, 1,
tzinfo=pendulum.timezone("UTC")),
+ table_names=["asset_event", "task_reschedule", "deadline"],
+ dag_ids=dag_ids,
+ exclude_dag_ids=exclude_dag_ids,
+ dry_run=False,
+ confirm=False,
+ error_on_cleanup_failure=True,
+ session=session,
+ )
+
+ @pytest.mark.parametrize(
+ "table_names",
+ [
+ pytest.param(["asset_event", "task_reschedule"],
id="child_tables_only"),
+ pytest.param(
+ ["asset_event", "task_reschedule", "task_instance", "dag_run"],
+ id="with_parents_that_cascade",
+ ),
+ ],
+ )
+ def test_cleanup_dag_filtering_keeps_rows_of_other_dags(self, table_names):
+ 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()
+
+ for asset_id, dag_id in enumerate(["dag1", "dag2"], start=1):
+ 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_version = DagVersion.get_latest_version(dag_id)
+
+ dag_run = DagRun(
+ dag_id,
+ run_id=f"{dag_id}_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=dag_version.id,
+ )
+ ti.dag_id = dag_id
+ ti.start_date = base_date
+ session.add(dag_run)
+ session.add(ti)
+ session.flush()
+
+ session.add(
+ TaskReschedule(
+ ti_id=ti.id,
+ start_date=base_date,
+ end_date=base_date.add(minutes=1),
+ reschedule_date=base_date.add(minutes=5),
+ )
+ )
+ session.add(
+ AssetEvent(asset_id=asset_id, source_dag_id=dag_id,
extra={}, timestamp=base_date)
+ )
+ session.commit()
+
+ run_cleanup(
+ clean_before_timestamp=base_date.add(days=10),
+ table_names=table_names,
+ dag_ids=["dag1"],
Review Comment:
**[P2] Populated test misses the new behavior**
This test hard-codes `dag_ids`, creates no Deadline rows, and never
exercises `exclude_dag_ids` with populated tables. Seed two Dags plus a
Deadline whose `dagrun_id` is NULL, test both inclusion and exclusion, and
assert the survivors. That would also expose the nullable
`AssetEvent.source_dag_id` defect.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/tests/unit/utils/test_db_cleanup.py:
##########
@@ -897,6 +1091,23 @@ def test_no_models_missing(self):
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:
Review Comment:
**[P2] Schema guard skips both relationship configurations**
This `continue` skips the two new `dag_id_via` entries entirely.
Consequently the regression guard validates neither the child foreign-key
column nor the parent primary-key and `dag_id` columns required by the cleanup
query. Extend the assertion to inspect both sides of each configured
relationship.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/tests/unit/utils/test_db_cleanup.py:
##########
@@ -638,7 +832,7 @@ def test_do_delete_success_does_not_call_rollback(self,
skip_archive, expected_c
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)]
Review Comment:
**[P3] Avoid an unrestricted MagicMock**
Airflow's review rules require mocks to use `spec` or `spec_set`. A concrete
result stub such as `SimpleNamespace(rowcount=1)`, or a
SQLAlchemy-result-specced mock, expresses the dependency more accurately and
prevents imaginary attributes.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
##########
airflow-core/tests/unit/utils/test_db_cleanup.py:
##########
@@ -551,6 +668,83 @@ def
test_dag_version_cleanup_skips_versions_pinned_by_task_instance(self):
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.
+
+ Simulates a race where a dag_version row passes the SELECT filter (no
TI
+ references it at archive-creation time) but a TI referencing it is
inserted
+ before the DELETE runs. The skip_if_referenced guard on the DELETE
itself
+ must leave the row in place instead of failing with IntegrityError.
+ """
+ from airflow.utils.db import reflect_tables
Review Comment:
**[P3] Move this import to module scope**
This function-local import has no documented circular-import or lazy-loading
requirement. Move `reflect_tables` to the module imports to follow the
repository's import rules.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
--
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]
