Re-anchor custom queued deadline references when a Dag run is cleared

airflow-core/newsfragments/71850.bugfix.rst

@@ -0,0 +1 @@
+Custom deadline references registered as ``DeadlineReference.TYPES.DAGRUN_QUEUED`` are now re-anchored to the new queued time when a Dag run is cleared, as the built-in ``DeadlineReference.DAGRUN_QUEUED_AT`` already was. Deadline alerts serialized before this fix pick the behaviour up once their Dag is parsed again.

airflow-core/src/airflow/models/deadline.py

@@ -58,6 +58,11 @@
 
 CALLBACK_METRICS_PREFIX = "deadline_alerts"
 
+# Must be in sync with the identically named constants in airflow.sdk.definitions.deadline
+EVALUATION_TIMING_FIELD = "evaluation_timing"
+DAGRUN_CREATED_TIMING = "DAGRUN_CREATED"
+DAGRUN_QUEUED_TIMING = "DAGRUN_QUEUED"
+
 
 class DeadlineDagRunProtocol(Protocol):
     """The Dag run interface deadline references are evaluated against."""
@@ -357,6 +362,8 @@ def get_reference_class(cls, reference_name: str) -> type[BaseDeadlineReference]
     class BaseDeadlineReference(LoggingMixin, ABC):
         """Base class for all Deadline implementations."""
 
+        evaluation_timing: str = DAGRUN_CREATED_TIMING
+
         @classproperty
         def reference_name(cls: Any) -> str:
             return cls.__name__
@@ -424,6 +431,8 @@ def _evaluate_with(self, *, session: Session, dagrun: DeadlineDagRunProtocol) ->
     class DagRunQueuedAtDeadline(BaseDeadlineReference):
         """A deadline that returns when a DagRun was queued."""
 
+        evaluation_timing: str = DAGRUN_QUEUED_TIMING
+
         def _evaluate_with(self, *, session: Session, dagrun: DeadlineDagRunProtocol) -> datetime | None:
             return dagrun.queued_at
 

airflow-core/src/airflow/models/serialized_dag.py

@@ -442,7 +442,7 @@ def _try_reuse_deadline_uuids(
         existing_deadline_uuids: list[str],
         new_deadline_data: list[dict],
         session: Session,
-    ) -> tuple[dict[str, dict], dict[str, str | None]] | None:
+    ) -> tuple[dict[str, dict], dict[str, dict[str, Any]]] | None:
         """
         Try to reuse existing deadline UUIDs if the deadline definitions haven't changed.
 
@@ -451,11 +451,11 @@ def _try_reuse_deadline_uuids(
         :param existing_deadline_uuids: List of UUID strings from existing serialized Dag
         :param new_deadline_data: List of new deadline alert data dicts from the Dag
         :param session: Database session
-        :return: Tuple of (uuid_mapping, name_updates) if all definitions match, None if any
+        :return: Tuple of (uuid_mapping, row_updates) if all definitions match, None if any
             mismatch detected.  ``uuid_mapping`` maps UUID string → new deadline data dict.
-            ``name_updates`` maps UUID string → new name **only** for entries whose name
-            changed relative to the existing DB row, so callers can issue targeted UPDATEs
-            and reliably detect whether any DB write occurred.
+            ``row_updates`` maps UUID string → the columns whose value **changed** relative to
+            the existing DB row, so callers can issue targeted UPDATEs and reliably detect
+            whether any DB write occurred.
         """
         # defensive check for old 3.1.x format
         if existing_deadline_uuids and not isinstance(existing_deadline_uuids[0], str):
@@ -464,8 +464,12 @@ def _try_reuse_deadline_uuids(
 
         def _definitions_match(deadline_data: dict, existing: DeadlineAlertModel) -> bool:
             """Check if raw deadline data matches an existing DeadlineAlert's definition."""
+            # The reference is compared on the fields the stored row already has: a field added to
+            # the serialized form by an upgrade would otherwise mark every stored alert as changed,
+            # orphaning the deadlines of running Dag runs. Such a row is refreshed in place below.
+            new_reference = deadline_data[DeadlineAlertFields.REFERENCE]
             return (
-                deadline_data[DeadlineAlertFields.REFERENCE] == existing.reference
+                all(new_reference.get(field) == value for field, value in existing.reference.items())
                 and deadline_data[DeadlineAlertFields.INTERVAL] == existing.interval
                 and deadline_data[DeadlineAlertFields.CALLBACK] == existing.callback_def
             )
@@ -483,7 +487,7 @@ def _definitions_match(deadline_data: dict, existing: DeadlineAlertModel) -> boo
 
         matched_uuids: set[UUID] = set()
         uuid_mapping: dict[str, dict] = {}
-        name_updates: dict[str, str | None] = {}
+        row_updates: dict[str, dict[str, Any]] = {}
 
         for deadline_alert in new_deadline_data:
             deadline_data = deadline_alert.get(Encoding.VAR, deadline_alert)
@@ -499,9 +503,15 @@ def _definitions_match(deadline_data: dict, existing: DeadlineAlertModel) -> boo
                     uuid_mapping[uuid_str] = deadline_data
                     matched_uuids.add(existing_alert.id)
                     found_match = True
+                    updates: dict[str, Any] = {}
                     new_name = deadline_data.get(DeadlineAlertFields.NAME)
                     if new_name != existing_alert.name:
-                        name_updates[uuid_str] = new_name
+                        updates["name"] = new_name
+                    new_reference = deadline_data[DeadlineAlertFields.REFERENCE]
+                    if new_reference != existing_alert.reference:
+                        updates["reference"] = new_reference
+                    if updates:
+                        row_updates[uuid_str] = updates
                     break
 
             if not found_match:
@@ -510,7 +520,7 @@ def _definitions_match(deadline_data: dict, existing: DeadlineAlertModel) -> boo
                 # to another deadline), so partial reuse would risk stale cross-references.
                 return None
 
-        return uuid_mapping, name_updates
+        return uuid_mapping, row_updates
 
     @classmethod
     def _create_deadline_alert_records(
@@ -653,7 +663,7 @@ def write_dag(
         serialized_dag_hash = _prefetched.dag_hash
         dag_version = _prefetched.dag_version
 
-        name_updated = False
+        deadline_rows_updated = False
         reused_deadline_data: dict[str, dict] | None = None
         if dag.data.get("dag", {}).get("deadline"):
             # The deadline handling below rewrites data["dag"]["deadline"] from a list of
@@ -679,17 +689,17 @@ def write_dag(
                 )
 
                 if reuse_result is not None:
-                    deadline_uuid_mapping, name_updates = reuse_result
+                    deadline_uuid_mapping, row_updates = reuse_result
                     # All deadlines matched — reuse the UUIDs to preserve hash.
-                    # Only issue UPDATE statements for rows whose name actually changed to
+                    # Only issue UPDATE statements for rows whose columns actually changed to
                     # avoid unnecessary writes and to make the return value accurate.
-                    for uuid_str, new_name in name_updates.items():
+                    for uuid_str, values in row_updates.items():
                         session.execute(
                             update(DeadlineAlertModel)
                             .where(DeadlineAlertModel.id == UUID(uuid_str))
-                            .values(name=new_name)
+                            .values(**values)
                         )
-                    name_updated = bool(name_updates)
+                    deadline_rows_updated = bool(row_updates)
                     dag.data["dag"]["deadline"] = existing_deadline_uuids
                     reused_deadline_data = deadline_uuid_mapping
                     deadline_uuid_mapping = {}
@@ -722,8 +732,8 @@ def write_dag(
                     dag_source_code=dag_source_code,
                     session=session,
                 )
-            if name_updated or bundle_metadata_changed:
-                # A write occurred — a deadline alert name update and/or a bundle
+            if deadline_rows_updated or bundle_metadata_changed:
+                # A write occurred — a deadline alert row update and/or a bundle
                 # metadata refresh — so report True so callers know the DB changed.
                 return True
             log.debug("Serialized DAG (%s) is unchanged. Skipping writing to DB", dag.dag_id)

airflow-core/src/airflow/models/taskinstance.py

@@ -101,7 +101,12 @@
 from airflow.models.base import Base, StringID
 from airflow.models.dag_version import DagVersion
 from airflow.models.dagbag import DBDagBag
-from airflow.models.deadline import Deadline, ReferenceModels
+from airflow.models.deadline import (
+    DAGRUN_QUEUED_TIMING,
+    EVALUATION_TIMING_FIELD,
+    Deadline,
+    ReferenceModels,
+)
 from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
 from airflow.models.dynamic_region import (
     SENTINEL_REGION_ID,
@@ -270,8 +275,13 @@ def _recalculate_dagrun_queued_at_deadlines(dagrun: DagRun, *, session: Session)
         .where(
             Deadline.dagrun_id == dagrun.id,
             Deadline.missed == false(),
-            DeadlineAlertModel.reference[ReferenceModels.REFERENCE_TYPE_FIELD].as_string()
-            == ReferenceModels.DagRunQueuedAtDeadline.__name__,
+            or_(
+                DeadlineAlertModel.reference[EVALUATION_TIMING_FIELD].as_string() == DAGRUN_QUEUED_TIMING,
+                # Alerts serialized before the timing was persisted only identify the built-in
+                # queued reference by name.
+                DeadlineAlertModel.reference[ReferenceModels.REFERENCE_TYPE_FIELD].as_string()
+                == ReferenceModels.DagRunQueuedAtDeadline.__name__,
+            ),
         )
     ).all()
 

airflow-core/src/airflow/serialization/definitions/deadline.py

@@ -27,7 +27,12 @@
 from sqlalchemy import select
 
 from airflow._shared.timezones import timezone
-from airflow.models.deadline import _get_evaluation_kwargs, classproperty
+from airflow.models.deadline import (
+    DAGRUN_CREATED_TIMING,
+    DAGRUN_QUEUED_TIMING,
+    _get_evaluation_kwargs,
+    classproperty,
+)
 from airflow.models.variable import Variable
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import provide_session
@@ -97,6 +102,8 @@ def get_reference_class(cls, reference_name: str) -> type[SerializedBaseDeadline
     class SerializedBaseDeadlineReference(LoggingMixin, ABC):
         """Base class for all serialized Deadline implementations."""
 
+        evaluation_timing: str = DAGRUN_CREATED_TIMING
+
         @classproperty
         def reference_name(cls: Any) -> str:
             return cls.__name__
@@ -161,6 +168,8 @@ def _evaluate_with(self, *, session: Session, dagrun: DeadlineDagRunProtocol) ->
     class DagRunQueuedAtDeadline(SerializedBaseDeadlineReference):
         """A deadline that returns when a DagRun was queued."""
 
+        evaluation_timing: str = DAGRUN_QUEUED_TIMING
+
         def _evaluate_with(self, *, session: Session, dagrun: DeadlineDagRunProtocol) -> datetime | None:
             return dagrun.queued_at
 
@@ -261,6 +270,7 @@ class SerializedCustomReference(SerializedBaseDeadlineReference):
 
         def __init__(self, inner_ref):
             self.inner_ref = inner_ref
+            self.evaluation_timing = getattr(inner_ref, "evaluation_timing", DAGRUN_CREATED_TIMING)
 
         @property
         def reference_name(self) -> str:

airflow-core/src/airflow/serialization/encoders.py

@@ -26,6 +26,7 @@
 import pendulum
 
 from airflow._shared.module_loading import qualname
+from airflow.models.deadline import DAGRUN_CREATED_TIMING, EVALUATION_TIMING_FIELD
 from airflow.partition_mappers.base import PartitionMapper as CorePartitionMapper
 from airflow.partition_mappers.wait_policy import WaitPolicy as CoreWaitPolicy
 from airflow.partition_mappers.window import Window as CoreWindow
@@ -283,6 +284,10 @@ def encode_deadline_reference(ref) -> dict[str, Any]:
 
     serialized = ref.serialize_reference()
 
+    # Added here rather than in serialize_reference() because custom references are expected to
+    # override that method, and an override would drop the timing.
+    serialized[EVALUATION_TIMING_FIELD] = getattr(ref, "evaluation_timing", DAGRUN_CREATED_TIMING)
+
     # Custom types (not built-in) need __class_path so the decoder can look them up.
     # Unlike built-in types which are looked up in SerializedReferenceModels,
     # custom types are resolved at deserialization time from the classes registered

airflow-core/tests/unit/models/test_deadline.py

@@ -34,6 +34,8 @@
 from airflow.sdk import timezone
 from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
 from airflow.sdk.definitions.deadline import (
+    DAGRUN_CREATED_TIMING,
+    DAGRUN_QUEUED_TIMING,
     AverageRuntimeDeadline,
     BaseDeadlineReference,
     DagRunLogicalDateDeadline,
@@ -627,14 +629,14 @@ def teardown_method(self):
         ],
     )
     @pytest.mark.parametrize(
-        "timing",
+        ("timing", "expected_evaluation_timing"),
         [
-            pytest.param(None, id="default_timing"),
-            pytest.param(DeadlineReference.TYPES.DAGRUN_CREATED, id="dagrun_created"),
-            pytest.param(DeadlineReference.TYPES.DAGRUN_QUEUED, id="dagrun_queued"),
+            pytest.param(None, DAGRUN_CREATED_TIMING, id="default_timing"),
+            pytest.param(DeadlineReference.TYPES.DAGRUN_CREATED, DAGRUN_CREATED_TIMING, id="dagrun_created"),
+            pytest.param(DeadlineReference.TYPES.DAGRUN_QUEUED, DAGRUN_QUEUED_TIMING, id="dagrun_queued"),
         ],
     )
-    def test_register_custom_reference(self, timing, reference):
+    def test_register_custom_reference(self, timing, expected_evaluation_timing, reference):
         if timing is None:
             result = DeadlineReference.register_custom_reference(reference)
             expected_timing = DeadlineReference.TYPES.DAGRUN_CREATED
@@ -645,6 +647,7 @@ def test_register_custom_reference(self, timing, reference):
         assert result is reference
         assert hasattr(DeadlineReference, reference.__name__)
         assert getattr(DeadlineReference, reference.__name__).__class__ is reference
+        assert reference.evaluation_timing == expected_evaluation_timing
 
         assert_correct_timing(reference, expected_timing)
         assert_builtin_types_unchanged(

airflow-core/tests/unit/models/test_deadline_alert.py

@@ -26,10 +26,20 @@
 
 from airflow._shared.timezones import timezone
 from airflow.exceptions import RemovedInAirflow4Warning
+from airflow.models import deadline as deadline_models
+from airflow.models.deadline import ReferenceModels
 from airflow.models.deadline_alert import DeadlineAlert
 from airflow.models.serialized_dag import SerializedDagModel
-from airflow.sdk.definitions.deadline import BaseDeadlineReference, DeadlineReference
+from airflow.sdk.definitions.deadline import (
+    DAGRUN_CREATED_TIMING,
+    DAGRUN_QUEUED_TIMING,
+    EVALUATION_TIMING_FIELD,
+    REFERENCE_TYPE_FIELD,
+    BaseDeadlineReference,
+    DeadlineReference,
+)
 from airflow.serialization.definitions.deadline import SerializedReferenceModels
+from airflow.serialization.encoders import encode_deadline_reference
 
 from tests_common.test_utils import db
 from unit.models import DEFAULT_DATE
@@ -41,6 +51,24 @@
 SERIALIZED_DAG_ID = "serialized_dag_uuid"
 
 
+class CustomQueuedReference(BaseDeadlineReference):
+    """A queued-anchored custom reference which overrides serialize_reference, as the docs suggest."""
+
+    evaluation_timing = DAGRUN_QUEUED_TIMING
+
+    def serialize_reference(self) -> dict:
+        return {REFERENCE_TYPE_FIELD: self.reference_name}
+
+
+def test_core_timing_constants_match_the_sdk():
+    """Core cannot import these from the SDK, so guard the copies against drift."""
+    assert (
+        deadline_models.EVALUATION_TIMING_FIELD,
+        deadline_models.DAGRUN_CREATED_TIMING,
+        deadline_models.DAGRUN_QUEUED_TIMING,
+    ) == (EVALUATION_TIMING_FIELD, DAGRUN_CREATED_TIMING, DAGRUN_QUEUED_TIMING)
+
+
 def _clean_db():
     db.clear_db_deadline_alert()
 
@@ -298,11 +326,32 @@ def _evaluate_with(self, *, session, **kwargs):
             "run_id": "test_run",
         }
 
+    @pytest.mark.parametrize(
+        ("reference", "expected_timing"),
+        [
+            pytest.param(DeadlineReference.DAGRUN_QUEUED_AT, DAGRUN_QUEUED_TIMING, id="sdk_queued"),
+            pytest.param(DeadlineReference.DAGRUN_LOGICAL_DATE, DAGRUN_CREATED_TIMING, id="sdk_created"),
+            pytest.param(CustomQueuedReference(), DAGRUN_QUEUED_TIMING, id="custom_queued"),
+            pytest.param(
+                SerializedReferenceModels.SerializedCustomReference(CustomQueuedReference()),
+                DAGRUN_QUEUED_TIMING,
+                id="wrapped_custom_queued",
+            ),
+            pytest.param(
+                SerializedReferenceModels.DagRunQueuedAtDeadline(),
+                DAGRUN_QUEUED_TIMING,
+                id="serialized_queued",
+            ),
+            pytest.param(
+                ReferenceModels.DagRunQueuedAtDeadline(), DAGRUN_QUEUED_TIMING, id="core_legacy_queued"
+            ),
+        ],
+    )
+    def test_encoded_reference_carries_evaluation_timing(self, reference, expected_timing):
+        assert encode_deadline_reference(reference)[EVALUATION_TIMING_FIELD] == expected_timing
+
     def test_core_deadline_reference_treated_as_builtins(self):
         """Test that refs from airflow.models.deadline are still treated as builtins."""
-        from airflow.models.deadline import ReferenceModels
-        from airflow.serialization.encoders import encode_deadline_reference
-
         ref = ReferenceModels.DagRunLogicalDateDeadline()
         serialized = encode_deadline_reference(ref)