From 94a9b2eb10a78689c40f3e5818e37d08e586b5a4 Mon Sep 17 00:00:00 2001 From: Shivam Rastogi <6463385+shivaam@users.noreply.github.com> Date: Wed, 24 Jun 2026 20:00:51 -0700 Subject: [PATCH 1/4] Handle missed deadlines when DagRuns fail --- airflow-core/src/airflow/models/dagrun.py | 38 +++++++++- airflow-core/tests/unit/models/test_dagrun.py | 70 +++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index edc4d903aa26c..41e08c1bebc3d 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -57,7 +57,16 @@ from sqlalchemy.ext.associationproxy import association_proxy from sqlalchemy.ext.hybrid import hybrid_property from sqlalchemy.ext.mutable import MutableDict -from sqlalchemy.orm import Mapped, declared_attr, joinedload, mapped_column, relationship, synonym, validates +from sqlalchemy.orm import ( + Mapped, + declared_attr, + joinedload, + mapped_column, + relationship, + selectinload, + synonym, + validates, +) from sqlalchemy.orm.exc import StaleDataError from sqlalchemy.sql.expression import false, select from sqlalchemy.sql.functions import coalesce @@ -1160,6 +1169,27 @@ def _emit_dagrun_span(self, state: DagRunState): span.set_status(status_code) span.end() + def _handle_missed_deadlines(self, *, session: Session) -> None: + deadline_query = ( + select(Deadline) + .where(Deadline.dagrun_id == self.id) + .where(~Deadline.missed) + .options(selectinload(Deadline.callback), selectinload(Deadline.dagrun)) + ) + try: + for deadline in session.scalars( + with_row_locks( + deadline_query, + of=Deadline, + session=session, + skip_locked=True, + key_share=False, + ) + ): + deadline.handle_miss(session) + except Exception: + self.log.warning("Failed to handle missed deadlines for %s", self, exc_info=True) + @provide_session def update_state( self, *, session: Session = NEW_SESSION, execute_callbacks: bool = True @@ -1260,6 +1290,9 @@ def recalculate(self) -> _UnfinishedStates: ) self._check_last_n_dagruns_failed(dag.dag_id, dag.max_consecutive_failed_dag_runs, session) + if dag.deadline: + self._handle_missed_deadlines(session=session) + # if all leaves succeeded and no unfinished tasks, the run succeeded elif not unfinished.tis and all(x.state in State.success_states for x in tis_for_dagrun_state): self.log.info("Marking run %s successful", self) @@ -1317,6 +1350,9 @@ def recalculate(self) -> _UnfinishedStates: execute=execute_callbacks, ) + if dag.deadline: + self._handle_missed_deadlines(session=session) + # finally, if the leaves aren't done, the dag is still running else: self.set_state(DagRunState.RUNNING) diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 4ec048a7f4bcc..4242ed72f54c9 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -1505,6 +1505,76 @@ 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, "handle_miss") + @mock.patch.object(Deadline, "prune_deadlines") + def test_dagrun_failure_handles_pending_deadline( + self, mock_prune, mock_handle_miss, session, deadline_test_dag + ): + scheduler_dag = deadline_test_dag( + deadline=DeadlineAlert( + reference=DeadlineReference.DAGRUN_QUEUED_AT, + interval=datetime.timedelta(hours=3), + callback=AsyncCallback(empty_callback_for_deadline), + ) + ) + + dag_run = self.create_dag_run( + dag=scheduler_dag, + task_states={"task_1": TaskInstanceState.SUCCESS, "task_2": TaskInstanceState.FAILED}, + session=session, + ) + dag_run.dag = scheduler_dag + session.add( + Deadline( + deadline_time=timezone.utcnow() + datetime.timedelta(hours=3), + callback=AsyncCallback(empty_callback_for_deadline), + dagrun_id=dag_run.id, + dag_id=dag_run.dag_id, + deadline_alert_id=None, + ) + ) + session.flush() + + dag_run.update_state(session=session) + + mock_handle_miss.assert_called_once() + mock_prune.assert_not_called() + assert dag_run.state == DagRunState.FAILED + + @mock.patch.object(Deadline, "handle_miss", side_effect=RuntimeError("deadline failure")) + def test_dagrun_failure_ignores_missed_deadline_handling_error( + self, mock_handle_miss, session, deadline_test_dag + ): + scheduler_dag = deadline_test_dag( + deadline=DeadlineAlert( + reference=DeadlineReference.DAGRUN_QUEUED_AT, + interval=datetime.timedelta(hours=3), + callback=AsyncCallback(empty_callback_for_deadline), + ) + ) + + dag_run = self.create_dag_run( + dag=scheduler_dag, + task_states={"task_1": TaskInstanceState.SUCCESS, "task_2": TaskInstanceState.FAILED}, + session=session, + ) + dag_run.dag = scheduler_dag + session.add( + Deadline( + deadline_time=timezone.utcnow() + datetime.timedelta(hours=3), + callback=AsyncCallback(empty_callback_for_deadline), + dagrun_id=dag_run.id, + dag_id=dag_run.dag_id, + deadline_alert_id=None, + ) + ) + session.flush() + + dag_run.update_state(session=session) + + mock_handle_miss.assert_called_once() + assert dag_run.state == DagRunState.FAILED + @mock.patch.object(Variable, "get") @mock.patch.object(Deadline, "prune_deadlines") def test_dagrun_deadline_variable_interval_stable(self, _, mock_get, session, deadline_test_dag): From 3163ee6e55e537dfd920d26c2e7960f3dbe21c01 Mon Sep 17 00:00:00 2001 From: Shivam Rastogi <6463385+shivaam@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:26:44 -0700 Subject: [PATCH 2/4] Add opt-in deadline alerts on DagRun failure --- airflow-core/docs/howto/deadline-alerts.rst | 1 + airflow-core/docs/migrations-ref.rst | 4 +- airflow-core/newsfragments/68961.feature.rst | 1 + .../core_api/datamodels/ui/deadline.py | 1 + .../core_api/openapi/_private_ui.yaml | 4 + ...0_add_fire_on_failure_to_deadline_alert.py | 48 +++++ airflow-core/src/airflow/models/dagrun.py | 43 ++-- .../src/airflow/models/deadline_alert.py | 5 +- .../src/airflow/models/serialized_dag.py | 2 + .../src/airflow/serialization/decoders.py | 1 + .../airflow/serialization/definitions/dag.py | 1 + .../serialization/definitions/deadline.py | 2 + .../src/airflow/serialization/encoders.py | 1 + .../ui/openapi-gen/requests/schemas.gen.ts | 8 +- .../ui/openapi-gen/requests/types.gen.ts | 3 +- .../core_api/routes/ui/test_deadlines.py | 1 + airflow-core/tests/unit/models/test_dagrun.py | 204 ++++++++++++++++-- .../tests/unit/models/test_deadline_alert.py | 12 ++ .../tests/unit/models/test_serialized_dag.py | 46 ++++ .../serialization/test_serialized_objects.py | 17 ++ .../src/airflow/sdk/definitions/deadline.py | 4 + .../task_sdk/definitions/test_deadline.py | 19 ++ 22 files changed, 395 insertions(+), 33 deletions(-) create mode 100644 airflow-core/newsfragments/68961.feature.rst create mode 100644 airflow-core/src/airflow/migrations/versions/0124_3_3_0_add_fire_on_failure_to_deadline_alert.py diff --git a/airflow-core/docs/howto/deadline-alerts.rst b/airflow-core/docs/howto/deadline-alerts.rst index 8e729516fefcb..5c426dc03153f 100644 --- a/airflow-core/docs/howto/deadline-alerts.rst +++ b/airflow-core/docs/howto/deadline-alerts.rst @@ -44,6 +44,7 @@ Creating a Deadline Alert requires three mandatory parameters: * Reference: When to start counting from * Interval: How far before or after the reference point to trigger the alert (either a timedelta or a dynamic interval such as VariableInterval) * Callback: A Callback object which contains a path to a callable and optional kwargs to pass to it if the deadline is exceeded +* Fire on failure: Optional. If set to ``True``, a pending Dag run deadline alert fires immediately when the scheduler marks the Dag run failed. Manual UI/API mark-failed paths are unchanged. The default is ``False``. Here is how Deadlines are calculated: diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index 3dd6e96a7d00b..ee107bb5729bd 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``c4e7a1f9b2d0`` (head) | ``436dc127462c`` | ``3.4.0`` | Add index on asset.uri. | +| ``b4f8c2a1d9e0`` (head) | ``c4e7a1f9b2d0`` | ``3.4.0`` | Add fire_on_failure to deadline_alert. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``c4e7a1f9b2d0`` | ``436dc127462c`` | ``3.4.0`` | Add index on asset.uri. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``436dc127462c`` | ``5a5d3253e946`` | ``3.4.0`` | Drop span_status column. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/newsfragments/68961.feature.rst b/airflow-core/newsfragments/68961.feature.rst new file mode 100644 index 0000000000000..05061953e5e16 --- /dev/null +++ b/airflow-core/newsfragments/68961.feature.rst @@ -0,0 +1 @@ +Add ``DeadlineAlert.fire_on_failure`` to optionally fire pending DagRun deadline alerts immediately when the scheduler marks a DagRun failed. diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/deadline.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/deadline.py index 6f9402f23603d..225be2730828d 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/deadline.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/deadline.py @@ -53,6 +53,7 @@ class DeadlineAlertResponse(BaseModel): name: str | None = None reference_type: str = Field(validation_alias=AliasPath("reference", "reference_type")) interval: float = Field(description="Interval in seconds between deadline evaluations.") + fire_on_failure: bool created_at: datetime diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml index e01a6f6e3c566..87fd946503aff 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml @@ -2851,6 +2851,9 @@ components: type: number title: Interval description: Interval in seconds between deadline evaluations. + fire_on_failure: + type: boolean + title: Fire On Failure created_at: type: string format: date-time @@ -2860,6 +2863,7 @@ components: - id - reference_type - interval + - fire_on_failure - created_at title: DeadlineAlertResponse description: DeadlineAlert serializer for responses. diff --git a/airflow-core/src/airflow/migrations/versions/0124_3_3_0_add_fire_on_failure_to_deadline_alert.py b/airflow-core/src/airflow/migrations/versions/0124_3_3_0_add_fire_on_failure_to_deadline_alert.py new file mode 100644 index 0000000000000..5dcb8612c5bf1 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0124_3_3_0_add_fire_on_failure_to_deadline_alert.py @@ -0,0 +1,48 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Add fire_on_failure to deadline_alert. + +Revision ID: b4f8c2a1d9e0 +Revises: c4e7a1f9b2d0 +Create Date: 2026-07-01 12:00:00.000000 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "b4f8c2a1d9e0" +down_revision = "c4e7a1f9b2d0" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Add fire_on_failure column to deadline_alert.""" + with op.batch_alter_table("deadline_alert", schema=None) as batch_op: + batch_op.add_column(sa.Column("fire_on_failure", sa.Boolean(), nullable=False, server_default="0")) + + +def downgrade(): + """Remove fire_on_failure column from deadline_alert.""" + with op.batch_alter_table("deadline_alert", schema=None) as batch_op: + batch_op.drop_column("fire_on_failure") diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 41e08c1bebc3d..6c758cec7d827 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -53,7 +53,7 @@ update, ) from sqlalchemy.dialects import postgresql -from sqlalchemy.exc import IntegrityError +from sqlalchemy.exc import DBAPIError, IntegrityError from sqlalchemy.ext.associationproxy import association_proxy from sqlalchemy.ext.hybrid import hybrid_property from sqlalchemy.ext.mutable import MutableDict @@ -1170,25 +1170,40 @@ def _emit_dagrun_span(self, state: DagRunState): span.end() def _handle_missed_deadlines(self, *, session: Session) -> None: + """Handle pending deadlines that opt in to firing when their DagRun fails.""" deadline_query = ( select(Deadline) + .join(DeadlineAlertModel, Deadline.deadline_alert_id == DeadlineAlertModel.id) .where(Deadline.dagrun_id == self.id) .where(~Deadline.missed) - .options(selectinload(Deadline.callback), selectinload(Deadline.dagrun)) + .where(DeadlineAlertModel.fire_on_failure.is_(True)) + .options( + selectinload(Deadline.callback), + selectinload(Deadline.dagrun), + selectinload(Deadline.deadline_alert), + ) ) - try: - for deadline in session.scalars( - with_row_locks( - deadline_query, - of=Deadline, - session=session, - skip_locked=True, - key_share=False, - ) - ): + for deadline in session.scalars( + with_row_locks( + deadline_query, + of=Deadline, + session=session, + skip_locked=True, + key_share=False, + ) + ): + deadline_id = deadline.id + try: deadline.handle_miss(session) - except Exception: - self.log.warning("Failed to handle missed deadlines for %s", self, exc_info=True) + except DBAPIError: + raise + except Exception: + self.log.warning( + "Failed to handle missed deadline %s for %s", + deadline_id, + self, + exc_info=True, + ) @provide_session def update_state( diff --git a/airflow-core/src/airflow/models/deadline_alert.py b/airflow-core/src/airflow/models/deadline_alert.py index 20bfed459eef9..0abb83db4716c 100644 --- a/airflow-core/src/airflow/models/deadline_alert.py +++ b/airflow-core/src/airflow/models/deadline_alert.py @@ -21,7 +21,7 @@ from uuid import UUID import uuid6 -from sqlalchemy import JSON, ForeignKey, String, Text, Uuid, select +from sqlalchemy import JSON, Boolean, ForeignKey, String, Text, Uuid, select from sqlalchemy.exc import NoResultFound from sqlalchemy.orm import Mapped, mapped_column @@ -52,6 +52,7 @@ class DeadlineAlert(Base): reference: Mapped[dict] = mapped_column(JSON, nullable=False) interval: Mapped[dict] = mapped_column(JSON, nullable=False) callback_def: Mapped[dict] = mapped_column(JSON, nullable=False) + fire_on_failure: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0") def __repr__(self): @@ -79,6 +80,7 @@ def __repr__(self): f"name={self.name or 'Unnamed'}, " f"reference={self.reference}, " f"interval={interval_display}, " + f"fire_on_failure={self.fire_on_failure}, " f"callback={self.callback_def}" ) @@ -90,6 +92,7 @@ def matches_definition(self, other: DeadlineAlert) -> bool: self.reference == other.reference and self.interval == other.interval and self.callback_def == other.callback_def + and self.fire_on_failure == other.fire_on_failure ) @property diff --git a/airflow-core/src/airflow/models/serialized_dag.py b/airflow-core/src/airflow/models/serialized_dag.py index fd773e8ed4fe3..19d90fc96e3d3 100644 --- a/airflow-core/src/airflow/models/serialized_dag.py +++ b/airflow-core/src/airflow/models/serialized_dag.py @@ -459,6 +459,7 @@ def _definitions_match(deadline_data: dict, existing: DeadlineAlertModel) -> boo deadline_data[DeadlineAlertFields.REFERENCE] == existing.reference and deadline_data[DeadlineAlertFields.INTERVAL] == existing.interval and deadline_data[DeadlineAlertFields.CALLBACK] == existing.callback_def + and deadline_data.get(DeadlineAlertFields.FIRE_ON_FAILURE, False) == existing.fire_on_failure ) if len(existing_deadline_uuids) != len(new_deadline_data): @@ -525,6 +526,7 @@ def _create_deadline_alert_records( reference=deadline_data[DeadlineAlertFields.REFERENCE], interval=deadline_data[DeadlineAlertFields.INTERVAL], callback_def=deadline_data[DeadlineAlertFields.CALLBACK], + fire_on_failure=deadline_data.get(DeadlineAlertFields.FIRE_ON_FAILURE, False), ) serialized_dag.deadline_alerts.append(alert) diff --git a/airflow-core/src/airflow/serialization/decoders.py b/airflow-core/src/airflow/serialization/decoders.py index 890a1296a4a00..be940922ba22e 100644 --- a/airflow-core/src/airflow/serialization/decoders.py +++ b/airflow-core/src/airflow/serialization/decoders.py @@ -196,6 +196,7 @@ def decode_deadline_alert(encoded_data: dict): interval=interval, callback=deserialize(data[DeadlineAlertFields.CALLBACK]), name=data.get(DeadlineAlertFields.NAME), + fire_on_failure=data.get(DeadlineAlertFields.FIRE_ON_FAILURE, False), ) diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py b/airflow-core/src/airflow/serialization/definitions/dag.py index abae0f9778572..8819fdc267024 100644 --- a/airflow-core/src/airflow/serialization/definitions/dag.py +++ b/airflow-core/src/airflow/serialization/definitions/dag.py @@ -724,6 +724,7 @@ def _process_dagrun_deadline_alerts( DeadlineAlertFields.REFERENCE: deadline_alert.reference, DeadlineAlertFields.INTERVAL: deadline_alert.interval, DeadlineAlertFields.CALLBACK: deadline_alert.callback_def, + DeadlineAlertFields.FIRE_ON_FAILURE: deadline_alert.fire_on_failure, }, } ) diff --git a/airflow-core/src/airflow/serialization/definitions/deadline.py b/airflow-core/src/airflow/serialization/definitions/deadline.py index 20fac2b54e874..7a20dad49bef9 100644 --- a/airflow-core/src/airflow/serialization/definitions/deadline.py +++ b/airflow-core/src/airflow/serialization/definitions/deadline.py @@ -55,6 +55,7 @@ class DeadlineAlertFields: REFERENCE = "reference" INTERVAL = "interval" CALLBACK = "callback" + FIRE_ON_FAILURE = "fire_on_failure" class SerializedReferenceModels: @@ -380,3 +381,4 @@ class SerializedDeadlineAlert: interval: timedelta | VariableInterval callback: Any name: str | None = None + fire_on_failure: bool = False diff --git a/airflow-core/src/airflow/serialization/encoders.py b/airflow-core/src/airflow/serialization/encoders.py index ab37b78550d77..acec187cbf43f 100644 --- a/airflow-core/src/airflow/serialization/encoders.py +++ b/airflow-core/src/airflow/serialization/encoders.py @@ -251,6 +251,7 @@ def encode_deadline_alert(d: DeadlineAlert | SerializedDeadlineAlert) -> dict[st "reference": encode_deadline_reference(d.reference), "interval": serialize(d.interval), "callback": serialize(d.callback), + "fire_on_failure": d.fire_on_failure, } diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index df738113acc1b..37ef750774b9f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -9426,6 +9426,10 @@ export const $DeadlineAlertResponse = { title: 'Interval', description: 'Interval in seconds between deadline evaluations.' }, + fire_on_failure: { + type: 'boolean', + title: 'Fire On Failure' + }, created_at: { type: 'string', format: 'date-time', @@ -9433,7 +9437,7 @@ export const $DeadlineAlertResponse = { } }, type: 'object', - required: ['id', 'reference_type', 'interval', 'created_at'], + required: ['id', 'reference_type', 'interval', 'fire_on_failure', 'created_at'], title: 'DeadlineAlertResponse', description: 'DeadlineAlert serializer for responses.' } as const; @@ -10968,4 +10972,4 @@ export const $UIAlert = { required: ['text', 'category'], title: 'UIAlert', description: 'Optional alert to be shown at the top of the page.' -} as const; \ No newline at end of file +} as const; diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 111ba9204ef1e..447400818c2f8 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2406,6 +2406,7 @@ export type DeadlineAlertResponse = { * Interval in seconds between deadline evaluations. */ interval: number; + fire_on_failure: boolean; created_at: string; }; @@ -8630,4 +8631,4 @@ export type $OpenApiTs = { }; }; }; -}; \ No newline at end of file +}; diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_deadlines.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_deadlines.py index acadab3a6b16f..4f64fb9c932b1 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_deadlines.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_deadlines.py @@ -489,6 +489,7 @@ def test_alert_response_fields(self, test_client): assert alert["name"] == ALERT_NAME assert alert["interval"] == 3600.0 assert alert["reference_type"] == "DagRunQueuedAtDeadline" + assert alert["fire_on_failure"] is False assert "id" in alert assert "created_at" in alert diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 4242ed72f54c9..ce36beb32ae11 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -40,6 +40,7 @@ select, update, ) +from sqlalchemy.exc import DBAPIError from sqlalchemy.orm import joinedload from sqlalchemy.orm.exc import StaleDataError @@ -71,7 +72,7 @@ task_group, teardown, ) -from airflow.sdk.definitions.callback import AsyncCallback +from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference, VariableInterval from airflow.sdk.definitions.variable import Variable from airflow.sdk.exceptions import AirflowRuntimeError @@ -110,6 +111,11 @@ async def empty_callback_for_deadline(): pass +def empty_sync_callback_for_deadline(): + """Used to confirm executor deadline callbacks can be queued without executing them.""" + pass + + @pytest.fixture(scope="module") def dagbag(): from airflow.dag_processing.dagbag import DagBag @@ -202,6 +208,27 @@ def create_dag_run( return dag_run + @staticmethod + def create_deadline_alert_model(dag: SerializedDAG, *, session: Session, fire_on_failure: bool): + serialized_dag = session.scalar( + select(SerializedDagModel) + .where(SerializedDagModel.dag_id == dag.dag_id) + .order_by(SerializedDagModel.created_at.desc()) + .limit(1) + ) + assert serialized_dag is not None + + deadline_alert = DeadlineAlertModel( + serialized_dag_id=serialized_dag.id, + reference=DeadlineReference.DAGRUN_QUEUED_AT.serialize_reference(), + interval=60, + callback_def={"path": "tests.unit.models.test_dagrun.empty_callback_for_deadline"}, + fire_on_failure=fire_on_failure, + ) + session.add(deadline_alert) + session.flush() + return deadline_alert + def test_clear_task_instances_for_backfill_running_dagrun(self, dag_maker, session): now = timezone.utcnow() state = DagRunState.RUNNING @@ -1510,19 +1537,94 @@ def test_dagrun_success_handles_empty_deadline_list(self, mock_prune, dag_maker, def test_dagrun_failure_handles_pending_deadline( self, mock_prune, mock_handle_miss, session, deadline_test_dag ): - scheduler_dag = deadline_test_dag( - deadline=DeadlineAlert( - reference=DeadlineReference.DAGRUN_QUEUED_AT, - interval=datetime.timedelta(hours=3), + scheduler_dag = deadline_test_dag() + + dag_run = self.create_dag_run( + dag=scheduler_dag, + task_states={"task_1": TaskInstanceState.SUCCESS, "task_2": TaskInstanceState.FAILED}, + session=session, + ) + deadline_alert = self.create_deadline_alert_model( + scheduler_dag, session=session, fire_on_failure=True + ) + scheduler_dag.deadline = [str(deadline_alert.id)] + dag_run.dag = scheduler_dag + session.add( + Deadline( + deadline_time=timezone.utcnow() + datetime.timedelta(hours=3), + callback=AsyncCallback(empty_callback_for_deadline), + dagrun_id=dag_run.id, + dag_id=dag_run.dag_id, + deadline_alert_id=deadline_alert.id, + ) + ) + session.flush() + + dag_run.update_state(session=session) + + mock_handle_miss.assert_called_once() + mock_prune.assert_not_called() + assert dag_run.state == DagRunState.FAILED + + @mock.patch.object(Deadline, "handle_miss") + def test_dagrun_deadlock_handles_pending_deadline(self, mock_handle_miss, dag_maker, session): + with dag_maker( + dag_id="test_dagrun_deadlock_handles_pending_deadline", + schedule=datetime.timedelta(days=1), + session=session, + ): + up = EmptyOperator(task_id="upstream") + middle = EmptyOperator(task_id="wrong") + down = EmptyOperator(task_id="downstream") + + middle.trigger_rule = TriggerRule.ONE_FAILED + middle.set_upstream(up) + middle.set_downstream(down) + + dag_run = dag_maker.create_dagrun() + scheduler_dag = dag_run.get_dag() + deadline_alert = self.create_deadline_alert_model( + scheduler_dag, session=session, fire_on_failure=True + ) + scheduler_dag.deadline = [str(deadline_alert.id)] + dag_run.dag = scheduler_dag + session.add( + Deadline( + deadline_time=timezone.utcnow() + datetime.timedelta(hours=3), callback=AsyncCallback(empty_callback_for_deadline), + dagrun_id=dag_run.id, + dag_id=dag_run.dag_id, + deadline_alert_id=deadline_alert.id, ) ) + ti_up: TI = dag_run.get_task_instance(task_id=up.task_id, session=session) + ti_middle: TI = dag_run.get_task_instance(task_id=middle.task_id, session=session) + ti_up.set_state(state=TaskInstanceState.SUCCESS, session=session) + ti_middle.set_state(state=None, session=session) + ti_middle.task.trigger_rule = "invalid" + session.flush() + + dag_run.update_state(session=session) + + mock_handle_miss.assert_called_once() + assert dag_run.state == DagRunState.FAILED + + @mock.patch.object(Deadline, "handle_miss") + def test_dagrun_failure_skips_pending_deadline_by_default( + self, mock_handle_miss, session, deadline_test_dag + ): + scheduler_dag = deadline_test_dag() + dag_run = self.create_dag_run( dag=scheduler_dag, task_states={"task_1": TaskInstanceState.SUCCESS, "task_2": TaskInstanceState.FAILED}, session=session, ) + deadline_alert = self.create_deadline_alert_model( + scheduler_dag, session=session, fire_on_failure=False + ) + scheduler_dag.deadline = [str(deadline_alert.id)] dag_run.dag = scheduler_dag session.add( Deadline( @@ -1530,34 +1632,67 @@ def test_dagrun_failure_handles_pending_deadline( callback=AsyncCallback(empty_callback_for_deadline), dagrun_id=dag_run.id, dag_id=dag_run.dag_id, - deadline_alert_id=None, + deadline_alert_id=deadline_alert.id, ) ) session.flush() dag_run.update_state(session=session) - mock_handle_miss.assert_called_once() - mock_prune.assert_not_called() + mock_handle_miss.assert_not_called() assert dag_run.state == DagRunState.FAILED @mock.patch.object(Deadline, "handle_miss", side_effect=RuntimeError("deadline failure")) def test_dagrun_failure_ignores_missed_deadline_handling_error( self, mock_handle_miss, session, deadline_test_dag ): - scheduler_dag = deadline_test_dag( - deadline=DeadlineAlert( - reference=DeadlineReference.DAGRUN_QUEUED_AT, - interval=datetime.timedelta(hours=3), + scheduler_dag = deadline_test_dag() + + dag_run = self.create_dag_run( + dag=scheduler_dag, + task_states={"task_1": TaskInstanceState.SUCCESS, "task_2": TaskInstanceState.FAILED}, + session=session, + ) + deadline_alert = self.create_deadline_alert_model( + scheduler_dag, session=session, fire_on_failure=True + ) + scheduler_dag.deadline = [str(deadline_alert.id)] + dag_run.dag = scheduler_dag + session.add( + Deadline( + deadline_time=timezone.utcnow() + datetime.timedelta(hours=3), callback=AsyncCallback(empty_callback_for_deadline), + dagrun_id=dag_run.id, + dag_id=dag_run.dag_id, + deadline_alert_id=deadline_alert.id, ) ) + session.flush() + + dag_run.update_state(session=session) + + mock_handle_miss.assert_called_once() + assert dag_run.state == DagRunState.FAILED + + @mock.patch.object( + Deadline, + "handle_miss", + side_effect=DBAPIError("statement", {}, Exception("deadline db failure")), + ) + def test_dagrun_failure_reraises_missed_deadline_db_error( + self, mock_handle_miss, session, deadline_test_dag + ): + scheduler_dag = deadline_test_dag() dag_run = self.create_dag_run( dag=scheduler_dag, task_states={"task_1": TaskInstanceState.SUCCESS, "task_2": TaskInstanceState.FAILED}, session=session, ) + deadline_alert = self.create_deadline_alert_model( + scheduler_dag, session=session, fire_on_failure=True + ) + scheduler_dag.deadline = [str(deadline_alert.id)] dag_run.dag = scheduler_dag session.add( Deadline( @@ -1565,15 +1700,56 @@ def test_dagrun_failure_ignores_missed_deadline_handling_error( callback=AsyncCallback(empty_callback_for_deadline), dagrun_id=dag_run.id, dag_id=dag_run.dag_id, - deadline_alert_id=None, + deadline_alert_id=deadline_alert.id, ) ) session.flush() - dag_run.update_state(session=session) + with pytest.raises(DBAPIError): + dag_run.update_state(session=session) mock_handle_miss.assert_called_once() + + def test_dagrun_failure_handles_pending_deadline_under_prohibit_commit(self, session, deadline_test_dag): + scheduler_dag = deadline_test_dag() + + dag_run = self.create_dag_run( + dag=scheduler_dag, + task_states={"task_1": TaskInstanceState.SUCCESS, "task_2": TaskInstanceState.FAILED}, + session=session, + ) + deadline_alert = self.create_deadline_alert_model( + scheduler_dag, session=session, fire_on_failure=True + ) + scheduler_dag.deadline = [str(deadline_alert.id)] + dag_run.dag = scheduler_dag + deadline = Deadline( + deadline_time=timezone.utcnow() + datetime.timedelta(hours=3), + callback=SyncCallback(empty_sync_callback_for_deadline), + dagrun_id=dag_run.id, + dag_id=dag_run.dag_id, + deadline_alert_id=deadline_alert.id, + ) + session.add(deadline) + session.flush() + + dag_run_id = dag_run.id + deadline_id = deadline.id + with prohibit_commit(session) as guard: + dag_run.update_state(session=session) + guard.commit() + + dag_run = session.get(DagRun, dag_run_id) + deadline = session.get(Deadline, deadline_id) + + assert dag_run is not None + assert deadline is not None assert dag_run.state == DagRunState.FAILED + assert deadline.missed is True + assert deadline.callback.state == "pending" + assert deadline.callback.data["dag_run_id"] == str(dag_run.id) + assert deadline.callback.data["dag_id"] == dag_run.dag_id + assert deadline.callback.data["deadline_id"] == str(deadline.id) @mock.patch.object(Variable, "get") @mock.patch.object(Deadline, "prune_deadlines") diff --git a/airflow-core/tests/unit/models/test_deadline_alert.py b/airflow-core/tests/unit/models/test_deadline_alert.py index a9b1854f6abce..b51ee990c2d34 100644 --- a/airflow-core/tests/unit/models/test_deadline_alert.py +++ b/airflow-core/tests/unit/models/test_deadline_alert.py @@ -84,6 +84,7 @@ def test_deadline_alert_creation(self, deadline_alert_orm): assert deadline_alert_orm.id is not None assert deadline_alert_orm.created_at == DEFAULT_DATE assert deadline_alert_orm.name == DEADLINE_NAME + assert deadline_alert_orm.fire_on_failure is False def test_minimal_deadline_alert_creation(self, dag_maker, session, deadline_reference): with dag_maker(DAG_ID, session=session): @@ -106,6 +107,7 @@ def test_minimal_deadline_alert_creation(self, dag_maker, session, deadline_refe assert deadline_alert.id is not None assert deadline_alert.created_at == DEFAULT_DATE assert deadline_alert.name is None + assert deadline_alert.fire_on_failure is False def test_deadline_alert_repr(self, deadline_alert_orm, deadline_reference): repr_str = repr(deadline_alert_orm) @@ -115,6 +117,7 @@ def test_deadline_alert_repr(self, deadline_alert_orm, deadline_reference): assert f"name={DEADLINE_NAME}" in repr_str assert f"reference={deadline_reference}" in repr_str assert "interval=1m" in repr_str + assert "fire_on_failure=False" in repr_str assert repr(deadline_alert_orm.callback_def) in repr_str def test_deadline_alert_matches_definition(self, session, deadline_reference): @@ -156,6 +159,15 @@ def test_deadline_alert_matches_definition(self, session, deadline_reference): ) assert not alert1.matches_definition(different_callback) + different_fire_on_failure = DeadlineAlert( + serialized_dag_id=SERIALIZED_DAG_ID, + reference=deadline_reference, + interval=DEADLINE_INTERVAL, + callback_def=DEADLINE_CALLBACK, + fire_on_failure=True, + ) + assert not alert1.matches_definition(different_fire_on_failure) + assert alert1.matches_definition("not a deadline alert") is NotImplemented def test_deadline_alert_reference_class_property(self, deadline_alert_orm): diff --git a/airflow-core/tests/unit/models/test_serialized_dag.py b/airflow-core/tests/unit/models/test_serialized_dag.py index 03a4159eb2b70..5b2408e62e011 100644 --- a/airflow-core/tests/unit/models/test_serialized_dag.py +++ b/airflow-core/tests/unit/models/test_serialized_dag.py @@ -981,6 +981,52 @@ def test_deadline_interval_change_triggers_new_serdag(self, testing_dag_bundle, assert new_serdag.dag_hash != orig_serdag.dag_hash assert new_alert.interval["__data__"] == 600.0 + def test_deadline_fire_on_failure_change_triggers_new_serdag(self, testing_dag_bundle, session): + dag_id = "test_fire_on_failure_change" + + 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) + scheduler_dag = sync_dag_to_db(dag, session=session) + scheduler_dag.create_dagrun( + run_id="test1", + run_after=DEFAULT_DATE, + state=DagRunState.QUEUED, + logical_date=DEFAULT_DATE, + data_interval=(DEFAULT_DATE, DEFAULT_DATE), + triggered_by=DagRunTriggeredByType.TEST, + run_type=DagRunType.MANUAL, + ) + session.commit() + + orig_serdag = session.scalar(select(SDM).where(SDM.dag_id == dag_id).order_by(SDM.created_at.desc())) + orig_alert = session.scalar(select(DAM).where(DAM.serialized_dag_id == orig_serdag.id)) + assert orig_alert.fire_on_failure is False + + dag.deadline = DeadlineAlert( + reference=DeadlineReference.DAGRUN_QUEUED_AT, + interval=timedelta(minutes=5), + callback=AsyncCallback(empty_callback_for_deadline), + fire_on_failure=True, + ) + + SDM.write_dag(LazyDeserializedDAG.from_dag(dag), bundle_name="testing", session=session) + session.commit() + + new_serdag_count = session.scalar(select(func.count()).select_from(SDM).where(SDM.dag_id == dag_id)) + new_serdag = session.scalar(select(SDM).where(SDM.dag_id == dag_id).order_by(SDM.created_at.desc())) + new_alert = session.scalar(select(DAM).where(DAM.serialized_dag_id == new_serdag.id)) + + assert new_serdag_count == 2 + assert new_serdag.dag_hash != orig_serdag.dag_hash + assert new_alert.fire_on_failure is True + def test_deadline_name_change_updates_db_and_returns_true(self, testing_dag_bundle, session): """Name-only deadline change: UUID reused, DB row updated, write_dag returns True.""" dag_id = "test_deadline_name_change" diff --git a/airflow-core/tests/unit/serialization/test_serialized_objects.py b/airflow-core/tests/unit/serialization/test_serialized_objects.py index 8495100cc47fc..4f747f84044e8 100644 --- a/airflow-core/tests/unit/serialization/test_serialized_objects.py +++ b/airflow-core/tests/unit/serialization/test_serialized_objects.py @@ -500,6 +500,7 @@ def test_serialize_deserialize_deadline_alert(reference): reference=reference, interval=timedelta(hours=1), callback=AsyncCallback(empty_callback_for_deadline, kwargs=TEST_CALLBACK_KWARGS), + fire_on_failure=True, ) # Use BaseSerialization like assets do @@ -511,6 +512,22 @@ def test_serialize_deserialize_deadline_alert(reference): assert deserialized.reference.serialize_reference() == reference.serialize_reference() assert deserialized.interval == original.interval assert deserialized.callback == original.callback + assert deserialized.fire_on_failure is True + + +def test_deserialize_deadline_alert_defaults_fire_on_failure_false(): + original = DeadlineAlert( + reference=DeadlineReference.DAGRUN_QUEUED_AT, + interval=timedelta(hours=1), + callback=AsyncCallback(TEST_CALLBACK_PATH, kwargs=TEST_CALLBACK_KWARGS), + ) + + serialized = BaseSerialization.serialize(original) + del serialized[Encoding.VAR][DeadlineAlertFields.FIRE_ON_FAILURE] + + deserialized = BaseSerialization.deserialize(serialized) + + assert deserialized.fire_on_failure is False def test_deserialize_deadline_alert_none_interval_raises(): diff --git a/task-sdk/src/airflow/sdk/definitions/deadline.py b/task-sdk/src/airflow/sdk/definitions/deadline.py index a9da3dd3d3ea2..c3d3e60e7de4c 100644 --- a/task-sdk/src/airflow/sdk/definitions/deadline.py +++ b/task-sdk/src/airflow/sdk/definitions/deadline.py @@ -150,10 +150,12 @@ def __init__( interval: timedelta | VariableInterval, callback: Callback, name: str | None = None, + fire_on_failure: bool = False, ): self.reference = reference self.interval = interval self.name = name + self.fire_on_failure = fire_on_failure if not isinstance(callback, (AsyncCallback, SyncCallback)): raise ValueError(f"Callbacks of type {type(callback).__name__} are not currently supported") @@ -166,6 +168,7 @@ def __eq__(self, other: object) -> bool: isinstance(self.reference, type(other.reference)) and self.interval == other.interval and self.callback == other.callback + and self.fire_on_failure == other.fire_on_failure ) def __hash__(self) -> int: @@ -174,6 +177,7 @@ def __hash__(self) -> int: type(self.reference).__name__, self.interval, self.callback, + self.fire_on_failure, ) ) diff --git a/task-sdk/tests/task_sdk/definitions/test_deadline.py b/task-sdk/tests/task_sdk/definitions/test_deadline.py index b104980e4c986..562e6e1562135 100644 --- a/task-sdk/tests/task_sdk/definitions/test_deadline.py +++ b/task-sdk/tests/task_sdk/definitions/test_deadline.py @@ -91,6 +91,16 @@ class TestDeadlineAlert: False, id="different_kwargs", ), + pytest.param( + DeadlineAlert( + reference=DeadlineReference.DAGRUN_QUEUED_AT, + interval=timedelta(hours=1), + callback=TEST_DEADLINE_CALLBACK, + fire_on_failure=True, + ), + False, + id="different_fire_on_failure", + ), pytest.param("not a DeadlineAlert", False, id="non_deadline_alert"), ], ) @@ -141,6 +151,15 @@ def test_deadline_alert_in_set(self): alert_set = {alert1, alert2} assert len(alert_set) == 1 + def test_deadline_alert_fire_on_failure_defaults_false(self): + alert = DeadlineAlert( + reference=DeadlineReference.DAGRUN_QUEUED_AT, + interval=timedelta(hours=1), + callback=TEST_DEADLINE_CALLBACK, + ) + + assert alert.fire_on_failure is False + @pytest.mark.parametrize( ("callback_class"), [ From d5644e9fd1d61f19e3d764a5b0402689a8b592f8 Mon Sep 17 00:00:00 2001 From: Shivam Rastogi <6463385+shivaam@users.noreply.github.com> Date: Thu, 2 Jul 2026 15:25:28 -0700 Subject: [PATCH 3/4] Refresh generated files after rebase --- airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts | 2 +- airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts | 2 +- airflow-core/src/airflow/utils/db.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 37ef750774b9f..c22e6b8c3cc93 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -10972,4 +10972,4 @@ export const $UIAlert = { required: ['text', 'category'], title: 'UIAlert', description: 'Optional alert to be shown at the top of the page.' -} as const; +} as const; \ No newline at end of file diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 447400818c2f8..9d8eebbfc6117 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -8631,4 +8631,4 @@ export type $OpenApiTs = { }; }; }; -}; +}; \ No newline at end of file diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index ce4ef036d20fd..684981605ca46 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "c4e7a1f9b2d0", + "3.4.0": "b4f8c2a1d9e0", } # Prefix used to identify tables holding data moved during migration. From c229f08ab2dd4edc4602edf03fe2cc77cc76c31a Mon Sep 17 00:00:00 2001 From: Shivam Rastogi <6463385+shivaam@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:59:37 -0700 Subject: [PATCH 4/4] Address deadline alert review feedback --- ..._add_fire_on_failure_to_deadline_alert.py} | 0 airflow-core/tests/unit/models/test_dagrun.py | 31 +++++++++++++++++++ 2 files changed, 31 insertions(+) rename airflow-core/src/airflow/migrations/versions/{0124_3_3_0_add_fire_on_failure_to_deadline_alert.py => 0127_3_4_0_add_fire_on_failure_to_deadline_alert.py} (100%) diff --git a/airflow-core/src/airflow/migrations/versions/0124_3_3_0_add_fire_on_failure_to_deadline_alert.py b/airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_fire_on_failure_to_deadline_alert.py similarity index 100% rename from airflow-core/src/airflow/migrations/versions/0124_3_3_0_add_fire_on_failure_to_deadline_alert.py rename to airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_fire_on_failure_to_deadline_alert.py diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index ce36beb32ae11..5e22706a990a0 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -1566,6 +1566,37 @@ def test_dagrun_failure_handles_pending_deadline( mock_prune.assert_not_called() assert dag_run.state == DagRunState.FAILED + @mock.patch.object(Deadline, "handle_miss") + def test_dagrun_failure_skips_already_missed_deadline(self, mock_handle_miss, session, deadline_test_dag): + scheduler_dag = deadline_test_dag() + + dag_run = self.create_dag_run( + dag=scheduler_dag, + task_states={"task_1": TaskInstanceState.SUCCESS, "task_2": TaskInstanceState.FAILED}, + session=session, + ) + deadline_alert = self.create_deadline_alert_model( + scheduler_dag, session=session, fire_on_failure=True + ) + scheduler_dag.deadline = [str(deadline_alert.id)] + dag_run.dag = scheduler_dag + deadline = Deadline( + deadline_time=timezone.utcnow() - datetime.timedelta(minutes=1), + callback=AsyncCallback(empty_callback_for_deadline), + dagrun_id=dag_run.id, + dag_id=dag_run.dag_id, + deadline_alert_id=deadline_alert.id, + ) + deadline.missed = True + session.add(deadline) + session.flush() + + dag_run.update_state(session=session) + + mock_handle_miss.assert_not_called() + assert deadline.missed is True + assert dag_run.state == DagRunState.FAILED + @mock.patch.object(Deadline, "handle_miss") def test_dagrun_deadlock_handles_pending_deadline(self, mock_handle_miss, dag_maker, session): with dag_maker(