Copilot commented on code in PR #62918:
URL: https://github.com/apache/airflow/pull/62918#discussion_r3066484763
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -1360,6 +1360,36 @@ def
test_dagrun_success_handles_empty_deadline_list(self, mock_prune, dag_maker,
mock_prune.assert_not_called()
assert dag_run.state == DagRunState.SUCCESS
+ @mock.patch.object(Deadline, "prune_deadlines")
+ @mock.patch.object(DeadlineAlertModel, "get_by_id")
+ def test_dagrun_failure_prunes_dagrun_deadlines(
+ self, mock_get_by_id, mock_prune, session, deadline_test_dag
+ ):
+ """Deadlines should be pruned when a DAG run fails, not just on
success."""
+ mock_deadline_alert = mock.MagicMock()
Review Comment:
`mock.MagicMock()` is created without a `spec`/`autospec`, which can hide
attribute/typo mistakes in tests. Consider using a spec’d mock (or an
autospecced patch) matching the `DeadlineAlert` model so the test fails if it
relies on non-existent attributes.
```suggestion
mock_deadline_alert = mock.create_autospec(DeadlineAlertModel,
instance=True)
```
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1267,6 +1255,21 @@ def recalculate(self) -> _UnfinishedStates:
self.data_interval_start,
self.data_interval_end,
)
+
+
+ if dag.deadline:
+ # The dagrun has reached a terminal state. Prune any pending
deadlines
+ # so they don't fire after the run is already finished.
+ deadline_alerts = [
+ DeadlineAlertModel.get_by_id(alert_id, session) for
alert_id in dag.deadline
+ ]
+
+ if any(
+ deadline_alert.reference_class in
SerializedReferenceModels.TYPES.DAGRUN
+ for deadline_alert in deadline_alerts
Review Comment:
`deadline_alerts = [DeadlineAlertModel.get_by_id(...)]` issues one DB query
per deadline ID. Since this code runs on every terminal DagRun (now including
failures), consider fetching all referenced DeadlineAlert rows in a single
query (`id IN (...)`) and then checking whether any reference_class is
DagRun-related, to avoid N queries when multiple deadlines are configured.
```suggestion
deadline_alert_reference_classes = session.execute(
select(DeadlineAlertModel.reference_class).where(
DeadlineAlertModel.id.in_(set(dag.deadline))
)
).scalars()
if any(
reference_class in SerializedReferenceModels.TYPES.DAGRUN
for reference_class in deadline_alert_reference_classes
```
--
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]
