Nataneljpwd commented on code in PR #75562:
URL: https://github.com/apache/airflow/pull/75562#discussion_r4434562918
##########
airflow-core/tests/unit/models/test_trigger.py:
##########
@@ -180,6 +180,96 @@ def test_clean_unused(session, dag_maker):
assert {result.id for result in results} == {trigger1.id, trigger4.id,
trigger5.id, trigger6.id}
+def _create_unused_trigger(session) -> int:
+ trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={})
+ session.add(trigger)
+ session.commit()
+ return trigger.id
+
+
+def _clean_unused_while_adding_reference(session, trigger_id, add_reference):
+ """
+ Run clean_unused, executing ``add_reference`` right before its DELETE
(#72061).
+
+ The trigger is unreferenced until then, so it surviving the cleanup also
shows the reference was added.
+ """
+
+ def add_reference_before_delete(connection, cursor, statement, parameters,
context, executemany):
+ normalized_statement = " ".join(statement.lower().replace('"',
"").replace("`", "").split())
+ if "delete from trigger" in normalized_statement:
+ connection.execute(add_reference)
+
+ event.listen(session.bind, "before_cursor_execute",
add_reference_before_delete)
+ try:
+ Trigger.clean_unused(session=session)
+ session.commit()
+ finally:
+ event.remove(session.bind, "before_cursor_execute",
add_reference_before_delete)
+
+ assert
session.scalar(select(func.count()).select_from(Trigger).where(Trigger.id ==
trigger_id)) == 1
+
+
+def test_clean_unused_keeps_trigger_deferred_onto_during_cleanup(session,
create_task_instance):
+ """A task deferring onto a trigger after the unused ids are read must not
be deleted."""
+ trigger_id = _create_unused_trigger(session)
+ task_instance = create_task_instance(
+ session=session, logical_date=timezone.utcnow(), state=State.DEFERRED
+ )
+ session.commit()
+
+ _clean_unused_while_adding_reference(
+ session,
+ trigger_id,
+ update(TaskInstance).where(TaskInstance.id ==
task_instance.id).values(trigger_id=trigger_id),
+ )
+
+ assert (
+ session.scalar(
+
select(func.count()).select_from(TaskInstance).where(TaskInstance.id ==
task_instance.id)
+ )
+ == 1
+ )
Review Comment:
For people unfamiliar with the on delete cascade in the trigger this check
might be confusing, consider adding a comment
--
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]
