diff --git a/airflow-core/newsfragments/71850.bugfix.rst b/airflow-core/newsfragments/71850.bugfix.rst new file mode 100644 index 0000000000000..a8e5ebb341421 --- /dev/null +++ b/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. diff --git a/airflow-core/src/airflow/models/deadline.py b/airflow-core/src/airflow/models/deadline.py index b9c212dc834e9..1fb1030cc852b 100644 --- a/airflow-core/src/airflow/models/deadline.py +++ b/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 diff --git a/airflow-core/src/airflow/models/serialized_dag.py b/airflow-core/src/airflow/models/serialized_dag.py index c8fcea2d634bc..e74ee1f9370d1 100644 --- a/airflow-core/src/airflow/models/serialized_dag.py +++ b/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) diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index c90deef1b60f7..68b921d04851b 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/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() diff --git a/airflow-core/src/airflow/serialization/definitions/deadline.py b/airflow-core/src/airflow/serialization/definitions/deadline.py index c7f1f1ee1119d..3522e2f378fad 100644 --- a/airflow-core/src/airflow/serialization/definitions/deadline.py +++ b/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: diff --git a/airflow-core/src/airflow/serialization/encoders.py b/airflow-core/src/airflow/serialization/encoders.py index d1f801d0b5337..dd683fa5b00f5 100644 --- a/airflow-core/src/airflow/serialization/encoders.py +++ b/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 diff --git a/airflow-core/tests/unit/models/test_deadline.py b/airflow-core/tests/unit/models/test_deadline.py index 6c09ff069b1f4..1d93ead8f7ff7 100644 --- a/airflow-core/tests/unit/models/test_deadline.py +++ b/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( diff --git a/airflow-core/tests/unit/models/test_deadline_alert.py b/airflow-core/tests/unit/models/test_deadline_alert.py index 54a138b64a077..90b85586d2dd6 100644 --- a/airflow-core/tests/unit/models/test_deadline_alert.py +++ b/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) diff --git a/airflow-core/tests/unit/models/test_serialized_dag.py b/airflow-core/tests/unit/models/test_serialized_dag.py index 7b47f70fdf2ba..7bec095c89062 100644 --- a/airflow-core/tests/unit/models/test_serialized_dag.py +++ b/airflow-core/tests/unit/models/test_serialized_dag.py @@ -42,7 +42,12 @@ from airflow.providers.standard.operators.python import PythonOperator from airflow.sdk import DAG, Asset, AssetAlias, task as task_decorator from airflow.sdk.definitions.callback import AsyncCallback -from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference +from airflow.sdk.definitions.deadline import ( + DAGRUN_QUEUED_TIMING, + EVALUATION_TIMING_FIELD, + DeadlineAlert, + DeadlineReference, +) from airflow.serialization.dag_dependency import DagDependency from airflow.serialization.definitions.dag import SerializedDAG from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG @@ -1371,6 +1376,63 @@ def test_write_dag_with_deadline_passes_schema_validation(self, testing_dag_bund # Must not raise: the stored UUID-reference form has to satisfy the serialized Dag schema. DagSerialization.validate_schema(result.data) + def test_write_dag_persists_evaluation_timing(self, testing_dag_bundle, session): + """The stored alert must record when the reference is anchored, not just its class name.""" + dag_id = "test_deadline_evaluation_timing" + dag = DAG( + dag_id=dag_id, + deadline=DeadlineAlert( + reference=DeadlineReference.DAGRUN_QUEUED_AT, + interval=timedelta(minutes=5), + callback=AsyncCallback(empty_callback_for_deadline), + ), + ) + EmptyOperator(task_id="task1", dag=dag) + sync_dag_to_db(dag, session=session) + session.commit() + + serialized_dag = session.scalar(select(SDM).where(SDM.dag_id == dag_id)) + alert = session.scalar(select(DAM).where(DAM.serialized_dag_id == serialized_dag.id)) + assert alert.reference[EVALUATION_TIMING_FIELD] == DAGRUN_QUEUED_TIMING + + def test_write_dag_refreshes_alert_predating_a_serialized_reference_field( + self, testing_dag_bundle, session + ): + """An alert stored before a reference field existed keeps its row and gains the field.""" + dag_id = "test_deadline_reference_field_added" + dag = DAG( + dag_id=dag_id, + deadline=DeadlineAlert( + reference=DeadlineReference.DAGRUN_QUEUED_AT, + interval=timedelta(minutes=5), + callback=AsyncCallback(empty_callback_for_deadline), + ), + ) + EmptyOperator(task_id="task1", dag=dag) + sync_dag_to_db(dag, session=session) + session.commit() + + def _count_versions() -> int: + return session.scalar( + select(func.count()).select_from(DagVersion).where(DagVersion.dag_id == dag_id) + ) + + alert = session.scalar(select(DAM)) + alert_id = alert.id + versions_before = _count_versions() + + # Mimic a row written by a version that did not serialize the evaluation timing yet. + alert.reference = {k: v for k, v in alert.reference.items() if k != EVALUATION_TIMING_FIELD} + session.commit() + + sync_dag_to_db(dag, session=session) + session.commit() + + alerts = session.scalars(select(DAM)).all() + assert [stored.id for stored in alerts] == [alert_id] + assert alerts[0].reference[EVALUATION_TIMING_FIELD] == DAGRUN_QUEUED_TIMING + assert _count_versions() == versions_before + def test_write_dag_does_not_mutate_caller_deadline_data(self, testing_dag_bundle, session): """write_dag must not rewrite the caller's LazyDeserializedDAG deadline in place. diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 4a7721b3a4117..e53032fc4c027 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -41,7 +41,7 @@ from sqlalchemy.orm import Session, load_only from sqlalchemy.orm.attributes import set_committed_value -from airflow import settings +from airflow import plugins_manager, settings from airflow._shared.observability.metrics.base_stats_logger import StatsLogger from airflow._shared.observability.traces import new_dagrun_trace_carrier, new_task_run_carrier from airflow._shared.timezones import timezone @@ -97,7 +97,13 @@ from airflow.sdk.api.datamodels._generated import AssetEventResponse, AssetResponse from airflow.sdk.definitions._internal.loop import create_loop from airflow.sdk.definitions.callback import AsyncCallback -from airflow.sdk.definitions.deadline import DeadlineReference, VariableInterval +from airflow.sdk.definitions.deadline import ( + DAGRUN_CREATED_TIMING, + DAGRUN_QUEUED_TIMING, + BaseDeadlineReference, + DeadlineReference, + VariableInterval, +) from airflow.sdk.definitions.param import process_params from airflow.sdk.definitions.taskgroup import TaskGroup from airflow.sdk.execution_time.comms import AssetEventsResult @@ -105,7 +111,7 @@ from airflow.serialization.definitions.assets import SerializedAsset from airflow.serialization.definitions.baseoperator import SerializedBaseOperator from airflow.serialization.definitions.dag import SerializedDAG -from airflow.serialization.encoders import ensure_serialized_asset +from airflow.serialization.encoders import encode_deadline_reference, ensure_serialized_asset from airflow.serialization.serialized_objects import OperatorSerialization, create_scheduler_operator from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_deps import REQUEUEABLE_DEPS, RUNNING_DEPS @@ -4984,6 +4990,86 @@ def test_clear_task_instances_skips_deadline_with_unresolvable_interval(dag_make assert args[0] == deadline.id +@pytest.mark.parametrize( + ("timing", "should_recalculate"), + [ + pytest.param(DAGRUN_QUEUED_TIMING, True, id="queued"), + pytest.param(DAGRUN_CREATED_TIMING, False, id="created"), + ], +) +def test_clear_task_instances_recalculates_custom_queued_deadlines( + dag_maker, session, monkeypatch, timing, should_recalculate +): + """Test that clearing tasks recalculates custom deadlines registered as queued-anchored.""" + # An offset of its own, so the new deadline can only come from the reference's own evaluation + # and not from the queued_at the built-in reference would return. + offset = datetime.timedelta(minutes=30) + + class CustomReference(BaseDeadlineReference): + evaluation_timing = timing + + def _evaluate_with(self, *, session, dagrun): + return dagrun.queued_at + offset + + reference = encode_deadline_reference(CustomReference()) + # Decoding the alert to resolve its interval goes through the plugin registry, which is where + # a custom reference has to be registered for the scheduler to resolve it at all. + monkeypatch.setattr( + plugins_manager, + "get_deadline_references_plugins", + lambda: {reference["__class_path"]: CustomReference}, + ) + + with dag_maker( + dag_id="test_recalculate_custom_deadlines", + schedule=datetime.timedelta(days=1), + ) as dag: + EmptyOperator(task_id="task_1") + + dag_run = dag_maker.create_dagrun() + ti = dag_run.get_task_instance("task_1", session=session) + ti.set_state(TaskInstanceState.SUCCESS, session=session) + + original_queued_at = timezone.utcnow() - datetime.timedelta(hours=2) + dag_run.queued_at = original_queued_at + session.flush() + + interval = datetime.timedelta(hours=1) + deadline_alert = DeadlineAlertModel( + serialized_dag_id=session.scalar( + select(SerializedDagModel.id).where(SerializedDagModel.dag_id == dag.dag_id) + ), + reference=reference, + interval=serialize(interval), + callback_def=serialize(AsyncCallback(empty_callback_for_deadline)), + ) + session.add(deadline_alert) + session.flush() + + original_deadline_time = original_queued_at + offset + interval + session.add( + Deadline( + dagrun_id=dag_run.id, + deadline_alert_id=deadline_alert.id, + deadline_time=original_deadline_time, + callback=AsyncCallback(empty_callback_for_deadline), + dag_id=dag_run.dag_id, + ) + ) + session.flush() + + tis = session.scalars(select(TI).where(TI.dag_id == dag.dag_id, TI.run_id == dag_run.run_id)).all() + clear_task_instances(tis, session) + + dag_run = session.scalar(select(DagRun).where(DagRun.id == dag_run.id)) + deadline = session.scalar(select(Deadline).where(Deadline.dagrun_id == dag_run.id)) + + if should_recalculate: + assert deadline.deadline_time == dag_run.queued_at + offset + interval + else: + assert deadline.deadline_time == original_deadline_time + + def test_get_dagrun_loaded_but_none_returns_dagrun(dag_maker, session): """ Test that `get_dagrun()` fetches `DagRun` from DB when the `dag_run` diff --git a/task-sdk/src/airflow/sdk/definitions/deadline.py b/task-sdk/src/airflow/sdk/definitions/deadline.py index 61dbe09b57bd0..b0915acc5b348 100644 --- a/task-sdk/src/airflow/sdk/definitions/deadline.py +++ b/task-sdk/src/airflow/sdk/definitions/deadline.py @@ -38,6 +38,13 @@ # Field name used in serialization - must be in sync with SerializedReferenceModels.REFERENCE_TYPE_FIELD REFERENCE_TYPE_FIELD = "reference_type" +# Field name holding the evaluation timing of a reference in its serialized form. +EVALUATION_TIMING_FIELD = "evaluation_timing" + +# Values stored under EVALUATION_TIMING_FIELD; they mirror the DeadlineReference.TYPES names. +DAGRUN_CREATED_TIMING = "DAGRUN_CREATED" +DAGRUN_QUEUED_TIMING = "DAGRUN_QUEUED" + class BaseDeadlineReference(ABC): """ @@ -53,6 +60,10 @@ class BaseDeadlineReference(ABC): attribute of an ``AirflowPlugin``; see :external:doc:`howto/deadline-alerts`. """ + # Set by ``register_custom_reference`` and persisted by the serializer, because the in-memory + # ``DeadlineReference.TYPES`` registry only exists in the process that ran the Dag file. + evaluation_timing: str = DAGRUN_CREATED_TIMING + @property def reference_name(self) -> str: """Return the class name as the reference identifier.""" @@ -91,6 +102,8 @@ class DagRunLogicalDateDeadline(BaseDeadlineReference): class DagRunQueuedAtDeadline(BaseDeadlineReference): """A deadline that returns when a DagRun was queued.""" + evaluation_timing: str = DAGRUN_QUEUED_TIMING + @dataclass class FixedDatetimeDeadline(BaseDeadlineReference): @@ -344,8 +357,10 @@ def register_custom_reference( # Add to appropriate deadline_reference_type classification if deadline_reference_type is cls.TYPES.DAGRUN_CREATED: cls.TYPES.DAGRUN_CREATED = cls.TYPES.DAGRUN_CREATED + (reference_class,) + reference_class.evaluation_timing = DAGRUN_CREATED_TIMING elif deadline_reference_type is cls.TYPES.DAGRUN_QUEUED: cls.TYPES.DAGRUN_QUEUED = cls.TYPES.DAGRUN_QUEUED + (reference_class,) + reference_class.evaluation_timing = DAGRUN_QUEUED_TIMING else: raise ValueError( f"Invalid deadline reference type {deadline_reference_type}; "