rjgoyln commented on code in PR #72062:
URL: https://github.com/apache/airflow/pull/72062#discussion_r3926265302
##########
airflow-core/src/airflow/models/trigger.py:
##########
@@ -258,10 +258,20 @@ def clean_unused(cls, *, session: Session = NEW_SESSION)
-> None:
)
ids = with_row_locks(ids, session, of=cls, skip_locked=True,
key_share=False)
if get_dialect_name(session) == "mysql":
- # MySQL doesn't support DELETE with JOIN, so we need to do it in
two steps
+ # MySQL doesn't support a DELETE whose subquery selects from the
target table,
+ # so materialize the ids first. The DELETE re-checks the reference
predicates:
+ # a task can defer onto one of these triggers in between, and
deleting it
+ # would cascade-delete the task instance row.
ids_list = list(session.scalars(ids).all())
session.execute(
-
delete(Trigger).where(Trigger.id.in_(ids_list)).execution_options(synchronize_session=False)
+ delete(Trigger)
+ .where(
+ Trigger.id.in_(ids_list),
+ ~cls.assets.any(),
+ ~cls.callback.has(),
+ ~cls.task_instance.has(),
Review Comment:
Could we extract these three predicates and reuse them in both the SELECT
and DELETE?
This would avoid duplicating the conditions and reduce the risk of adding a
new reference type to one side but forgetting the other.
##########
airflow-core/tests/unit/models/test_trigger.py:
##########
@@ -180,6 +180,48 @@ def test_clean_unused(session, dag_maker):
assert {result.id for result in results} == {trigger1.id, trigger4.id,
trigger5.id, trigger6.id}
+def
test_clean_unused_keeps_triggers_referenced_after_candidate_select(session,
create_task_instance):
+ """
+ Tests that the MySQL two-step delete does not remove a trigger that a task
+ instance deferred onto after the candidate ids were selected -- that would
+ cascade-delete the task instance row itself. The dialect is pinned to
+ "mysql" so the two-step branch runs on every backend.
+ """
+ trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={})
+ session.add(trigger)
+ session.flush()
+
+ task_instance = create_task_instance(
+ session=session, logical_date=timezone.utcnow(), state=State.SCHEDULED
+ )
+ session.flush()
+
+ real_scalars = session.scalars
+
+ def scalars_then_defer(*args, **kwargs):
+ # session.scalars() executes the SELECT eagerly, so the trigger is
+ # already on the kill list when the task instance defers onto it.
+ result = real_scalars(*args, **kwargs)
+ task_instance.state = State.DEFERRED
+ task_instance.trigger_id = trigger.id
+ session.flush()
+ return result
+
+ with (
+ patch("airflow.models.trigger.get_dialect_name", return_value="mysql"),
+ patch.object(session, "scalars", side_effect=scalars_then_defer),
Review Comment:
Nit: could we add autospec=True here and assert mock_scalars.call_count == 1?
This would make the mock stricter and explicitly verify that the candidate
IDs are materialized once before the concurrent defer is injected.
##########
airflow-core/tests/unit/models/test_trigger.py:
##########
@@ -180,6 +180,48 @@ def test_clean_unused(session, dag_maker):
assert {result.id for result in results} == {trigger1.id, trigger4.id,
trigger5.id, trigger6.id}
+def
test_clean_unused_keeps_triggers_referenced_after_candidate_select(session,
create_task_instance):
Review Comment:
Would it be worth parameterizing this test to also cover the Asset and
Callback references?
The current test reproduces the reported TaskInstance case, so this would
just improve coverage of the other two predicates.
--
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]
