ferruzzi commented on code in PR #66350:
URL: https://github.com/apache/airflow/pull/66350#discussion_r4039095321
##########
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:
+ break
Review Comment:
Your 11 Sep version had the sequence right, but it passed
`select(archive_table)`, which doesn't have a NOT EXISTS guard. That's why
`continue` hung. The calling code always runs a guarded query built by
`_build_query`, so a skipped row drops out of the next pass, but an unguarded
SELECT keeps returning that row forever. Switching to `_build_query` fixed the
hang, but now the first pass returns zero rows, the DELETE is never reached,
and the guard isn't actually hit int he test. It needs both: the real
`_build_query` query and the TI inserted mid-pass rather than up front. Try
this:
```python
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 configured keep_last per dag_id, so a lone
version is always the
# keep_last survivor and never eligible. 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
# 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):
"""Insert the referencing TI in the window the race needs.
_do_delete reflects both source and target immediately after
committing
the archive CTAS and immediately before building the DELETE,
so this call
site is the only seam between the two. The MySQL branch
reflects the
target alone earlier in the same pass, hence keying on the
two-table call.
If that call is ever moved or inlined, this test stops
reproducing the
race -- and would pass while testing nothing, so the
``raced`` assertion
below is not optional.
"""
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"
```
What was happening is that with a single `DagVersion`, the SELECT was
returning zero for two reasons. `dag_version` is `keep_last` per `dag_id`, so
a lone version is always the survivor and is excluded regardless of any TI.
That's why this version creates two. `latest` will take the `keep_last` slot
so `raced_old` comes up for deletion.
I've tested it on postgres and sqlite. It works as written and fails with
the `if skip_if_referenced:` block removed as expected, so I think it should be
good regardless of backend. I don't have a mysql backend handy, but the CI
will test that side.
--
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]
