From 710a936d6bdac09d5ca2724e74a9941797e6b3ed Mon Sep 17 00:00:00 2001 From: eugeneo Date: Mon, 17 Aug 2026 12:42:30 -0400 Subject: [PATCH 1/2] Add opt-in auto-clear of failed child tasks to TriggerDagRunOperator When TriggerDagRunOperator waits for a triggered child run (wait_for_completion=True) and that child fails, retrying the parent task previously dead-ended on DagRunAlreadyExists, or with reset_dag_run re-ran the entire child from scratch and discarded already-successful tasks. This adds an opt-in auto_clear_failed_tasks flag: on the parent task's retry it clears only the failed tasks (and their downstream) of the existing child run so the child resumes from its point of failure while preserving work that already succeeded. reset_dag_run keeps precedence when both are set. The failed-only clear is delivered server-side through a new Execution API version (2026-11-13); against an older core the operator raises rather than silently downgrading to a whole-run clear. Includes operator, task-SDK and execution-API unit tests, plus a regression test asserting the failed-only clear is reached when reset_dag_run=False (the gap that let the feature ship as a no-op). closes: #53402 --- .../execution_api/datamodels/dagrun.py | 6 + .../execution_api/routes/dag_runs.py | 15 +- .../execution_api/versions/__init__.py | 7 + .../execution_api/versions/v2026_11_13.py | 35 ++ .../execution_api/datamodels/test_dagrun.py | 39 ++ .../versions/head/test_capabilities.py | 36 ++ .../versions/v2026_11_13/__init__.py | 16 + .../versions/v2026_11_13/test_dag_runs.py | 81 +++ providers/standard/docs/changelog.rst | 10 + .../docs/operators/trigger_dag_run.rst | 40 ++ .../standard/operators/trigger_dagrun.py | 104 +++- .../standard/operators/test_trigger_dagrun.py | 563 ++++++++++++++++++ task-sdk/src/airflow/sdk/api/client.py | 16 +- .../airflow/sdk/api/datamodels/_generated.py | 11 + task-sdk/src/airflow/sdk/exceptions.py | 2 + .../src/airflow/sdk/execution_time/comms.py | 1 + .../sdk/execution_time/schema/schema.json | 2 +- .../airflow/sdk/execution_time/supervisor.py | 9 +- .../airflow/sdk/execution_time/task_runner.py | 1 + task-sdk/tests/task_sdk/api/test_client.py | 82 +++ .../task_sdk/execution_time/test_comms.py | 13 + .../execution_time/test_supervisor.py | 22 + .../execution_time/test_task_runner.py | 42 ++ 23 files changed, 1141 insertions(+), 12 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_11_13.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/datamodels/test_dagrun.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_capabilities.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_11_13/__init__.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_11_13/test_dag_runs.py diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dagrun.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dagrun.py index dd1f1aeb86665..c4524f16fd1aa 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dagrun.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dagrun.py @@ -35,6 +35,12 @@ class TriggerDAGRunPayload(StrictBaseModel): note: str | None = None +class ClearDagRunPayload(StrictBaseModel): + """Schema for Clear DAG Run API request.""" + + only_failed: bool = False + + class DagRunStateResponse(BaseModel): """Schema for DAG Run State response.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py index 213f01f4ffa5f..f17c5c370c7da 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py @@ -21,7 +21,7 @@ from typing import Annotated from cadwyn import VersionedAPIRouter -from fastapi import HTTPException, Query, status +from fastapi import Body, HTTPException, Query, status from sqlalchemy import func, select from sqlalchemy.exc import NoResultFound @@ -30,7 +30,11 @@ from airflow.api_fastapi.common.db.common import AsyncSessionDep, SessionDep from airflow.api_fastapi.common.types import UtcDateTime from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT -from airflow.api_fastapi.execution_api.datamodels.dagrun import DagRunStateResponse, TriggerDAGRunPayload +from airflow.api_fastapi.execution_api.datamodels.dagrun import ( + ClearDagRunPayload, + DagRunStateResponse, + TriggerDAGRunPayload, +) from airflow.api_fastapi.execution_api.datamodels.taskinstance import DagRun from airflow.api_fastapi.execution_api.datamodels.token import TIToken from airflow.api_fastapi.execution_api.security import ( @@ -192,6 +196,7 @@ def clear_dag_run( run_id: str, session: SessionDep, dag_bag: DagBagDep, + payload: ClearDagRunPayload = Body(default_factory=ClearDagRunPayload), ) -> None: """Clear a Dag run.""" dm = session.scalar(select(DagModel).where(~DagModel.is_stale, DagModel.dag_id == dag_id).limit(1)) @@ -221,7 +226,11 @@ def clear_dag_run( dag = get_dag_for_run(dag_bag, dag_run=dag_run, session=session) resolved_run_on_latest = resolve_run_on_latest_version(None, dag_id, session) - dag.clear(run_id=run_id, run_on_latest_version=resolved_run_on_latest) + dag.clear( + run_id=run_id, + only_failed=payload.only_failed, + run_on_latest_version=resolved_run_on_latest, + ) @router.get( diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index 8936306bbf5a6..c86e7331d8aa1 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -67,11 +67,18 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30_xcom_params import ( AddPreviousIterationToXComFilterParams, ) +from airflow.api_fastapi.execution_api.versions.v2026_11_13 import AddOnlyFailedToClearDagRunPayload from airflow.api_fastapi.execution_api.versions.v2027_02_28 import AddDagProcessorJobEndpoints +# Minimum Execution API version that honours the ``only_failed`` clear scope. Clients (e.g. the +# TriggerDagRunOperator failed-only path added in F1-T3) must gate on this and raise an explicit +# error when negotiating against an older core, rather than silently clearing the whole run. +MIN_VERSION_ONLY_FAILED_CLEAR = "2026-11-13" + bundle = VersionBundle( HeadVersion(), Version("2027-02-28", AddDagProcessorJobEndpoints), + Version("2026-11-13", AddOnlyFailedToClearDagRunPayload), Version( "2026-10-30", AddArgBindingsToTIRunContext, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_11_13.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_11_13.py new file mode 100644 index 0000000000000..c3c3088db6108 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_11_13.py @@ -0,0 +1,35 @@ +# 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. + +from __future__ import annotations + +from cadwyn import VersionChange, schema + +from airflow.api_fastapi.execution_api.datamodels.dagrun import ClearDagRunPayload + + +class AddOnlyFailedToClearDagRunPayload(VersionChange): + """Add the ``only_failed`` field to the clear-dag-run request payload.""" + + description = __doc__ + + # Request-body-only additive field: for older versions the field simply does not exist on + # ``ClearDagRunPayload`` (which forbids extra fields), so an older client omitting it clears the + # whole run, and one that sends it is rejected rather than silently downgraded to a whole-run clear. + instructions_to_migrate_to_previous_version = ( + schema(ClearDagRunPayload).field("only_failed").didnt_exist, + ) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/datamodels/test_dagrun.py b/airflow-core/tests/unit/api_fastapi/execution_api/datamodels/test_dagrun.py new file mode 100644 index 0000000000000..8de3f3376ecb4 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/datamodels/test_dagrun.py @@ -0,0 +1,39 @@ +# 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. + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from airflow.api_fastapi.execution_api.datamodels.dagrun import ClearDagRunPayload + + +class TestClearDagRunPayload: + def test_only_failed_defaults_to_false(self): + """Omitting only_failed yields False so the route performs a whole-run clear.""" + payload = ClearDagRunPayload() + assert payload.only_failed is False + + @pytest.mark.parametrize("value", [True, False]) + def test_only_failed_accepts_bool(self, value): + payload = ClearDagRunPayload(only_failed=value) + assert payload.only_failed is value + + def test_only_failed_rejects_non_bool(self): + with pytest.raises(ValidationError): + ClearDagRunPayload(only_failed="not-a-bool") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_capabilities.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_capabilities.py new file mode 100644 index 0000000000000..acfdf646e655d --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_capabilities.py @@ -0,0 +1,36 @@ +# 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. + +from __future__ import annotations + +from airflow.api_fastapi.execution_api.versions import ( + MIN_VERSION_ONLY_FAILED_CLEAR, + bundle, +) + + +def test_min_version_only_failed_clear_is_registered(): + """@AC-FR004-01 The advertised failed-only clear minimum maps to a registered API version.""" + registered = {version.value for version in bundle.versions} + assert MIN_VERSION_ONLY_FAILED_CLEAR in registered + + +def test_min_version_only_failed_clear_is_supported_by_head(): + """@AC-FR004-02 The current head version is at least the failed-only clear minimum.""" + # bundle.versions is ordered newest-first; the first entry is head/latest. + latest = bundle.versions[0].value + assert latest >= MIN_VERSION_ONLY_FAILED_CLEAR diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_11_13/__init__.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_11_13/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_11_13/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_11_13/test_dag_runs.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_11_13/test_dag_runs.py new file mode 100644 index 0000000000000..4c39d16feb20a --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_11_13/test_dag_runs.py @@ -0,0 +1,81 @@ +# 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. + +from __future__ import annotations + +from unittest import mock + +import pytest +from cadwyn import generate_versioned_models + +from airflow.api_fastapi.execution_api.datamodels.dagrun import ClearDagRunPayload +from airflow.api_fastapi.execution_api.versions import bundle +from airflow.serialization.definitions.dag import SerializedDAG +from airflow.utils.state import DagRunState + +from tests_common.test_utils.db import clear_db_runs + +# The version that introduces failed-only clear and the last one that predates it. +CURRENT_VERSION = "2026-11-13" +PREVIOUS_VERSION = "2026-10-30" + + +class TestClearDagRunPayloadVersioning: + """The ``only_failed`` field is versioned additively: present in the new version, absent before it.""" + + def test_field_present_in_current_version(self): + """@AC-FR004-01 The new version's ClearDagRunPayload exposes the additive only_failed field.""" + versioned = generate_versioned_models(bundle)[CURRENT_VERSION][ClearDagRunPayload] + assert "only_failed" in versioned.model_fields + + def test_field_negotiated_away_in_previous_version(self): + """@AC-FR004-02 The previous version negotiates only_failed away (field absent ⇒ whole-run).""" + versioned = generate_versioned_models(bundle)[PREVIOUS_VERSION][ClearDagRunPayload] + assert "only_failed" not in versioned.model_fields + + +@pytest.mark.db_test +class TestDagRunClearVersionNegotiation: + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + @pytest.fixture + def old_ver_client(self, client): + """Client pinned to the version that predates the ``only_failed`` clear field.""" + client.headers["Airflow-API-Version"] = PREVIOUS_VERSION + return client + + def test_old_version_no_body_clears_whole_run(self, old_ver_client, session, dag_maker): + """@AC-FR004-01 An old-version client sending no body clears the whole run (only_failed defaults False).""" + from airflow.providers.standard.operators.empty import EmptyOperator + + dag_id = "test_old_clear_no_body" + run_id = "test_run_id" + with dag_maker(dag_id=dag_id, session=session, serialized=True): + EmptyOperator(task_id="test_task") + dag_maker.create_dagrun(run_id=run_id, state=DagRunState.SUCCESS) + session.commit() + + with mock.patch.object(SerializedDAG, "clear", autospec=True) as mock_clear: + response = old_ver_client.post(f"/execution/dag-runs/{dag_id}/{run_id}/clear") + + assert response.status_code == 204 + mock_clear.assert_called_once() + assert mock_clear.call_args.kwargs["only_failed"] is False diff --git a/providers/standard/docs/changelog.rst b/providers/standard/docs/changelog.rst index 38039f486a9b2..f0f8b85a589e3 100644 --- a/providers/standard/docs/changelog.rst +++ b/providers/standard/docs/changelog.rst @@ -35,6 +35,16 @@ Changelog --------- +.. note:: + ``TriggerDagRunOperator`` gains an opt-in ``auto_clear_failed_tasks`` flag (default ``False``). + When enabled and the triggered child Dag run already exists in a failed state, only its failed + tasks and their downstream are cleared and re-run on the next execution of the operator (for + example on a retry), preserving already-succeeded tasks instead of resetting the whole run. The + feature is synchronous-only in this version (``wait_for_completion=True``) and is not supported + with ``deferrable=True``. ``reset_dag_run`` still takes precedence when both are set. On Airflow + 3.x the failed-only clear is delivered via the Execution API and requires a new enough core; the + operator raises ``NotImplementedError`` rather than silently falling back to a whole-run clear. + 1.20.1 ...... diff --git a/providers/standard/docs/operators/trigger_dag_run.rst b/providers/standard/docs/operators/trigger_dag_run.rst index 2d783248cb2fa..4a0522963070d 100644 --- a/providers/standard/docs/operators/trigger_dag_run.rst +++ b/providers/standard/docs/operators/trigger_dag_run.rst @@ -29,3 +29,43 @@ Use the :class:`~airflow.providers.standard.operators.trigger_dagrun.TriggerDagR :dedent: 4 :start-after: [START howto_operator_trigger_dagrun] :end-before: [END howto_operator_trigger_dagrun] + +Automatically clearing failed child tasks +----------------------------------------- + +By default, when the triggered child Dag run already exists, ``TriggerDagRunOperator`` either +raises, skips, or -- with ``reset_dag_run=True`` -- clears and re-runs the *whole* child run. The +opt-in ``auto_clear_failed_tasks`` flag offers a narrower behavior: when the existing child run is in +a failed state, only its failed tasks (and their downstream) are cleared and re-run, so +already-succeeded upstream tasks are preserved and not re-executed. + +The flag is off by default. To enable it, set ``auto_clear_failed_tasks=True`` and run the operator +synchronously (``wait_for_completion=True``). The failed-only clear happens on the next execution of +the operator -- for example on a task retry -- so bound the attempts with ``retries``: + +.. code-block:: python + + TriggerDagRunOperator( + task_id="run_child", + trigger_dag_id="child_dag", + wait_for_completion=True, # sync mode (v1 scope) + retries=2, # bounds the auto-clear + retry + auto_clear_failed_tasks=True, # on retry, clear failed + downstream of the failed child run + ) + +Caveats +~~~~~~~ + +* **Synchronous-only (v1).** ``auto_clear_failed_tasks`` applies when ``wait_for_completion=True`` and + is not supported with ``deferrable=True`` -- setting both raises ``ValueError``. +* **Cosmetic window.** Until the operator re-executes and triggers the clear, the child run may still + be shown as ``failed`` in the UI. +* **Task idempotency.** Previously-succeeded tasks are not re-run and the failed ones are, so the + triggered Dag's tasks should be idempotent; re-running a non-idempotent task may cause duplicate + side effects. +* **Precedence with ``reset_dag_run``.** If both are set, ``reset_dag_run`` wins and the whole run is + cleared; at most one clear is performed. +* **Airflow 3.x core-version requirement.** On Airflow 3.x the failed-only clear is delivered + server-side via the Execution API and requires a new enough core (Execution API version + ``2026-11-13`` or later). Against an older core the operator raises ``NotImplementedError`` rather + than silently falling back to a whole-run clear. diff --git a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py index 0a6f036b3de9b..2cf3eef7a81fb 100644 --- a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py +++ b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py @@ -60,6 +60,17 @@ XCOM_LOGICAL_DATE_ISO = "trigger_logical_date_iso" XCOM_RUN_ID = "trigger_run_id" +# Minimum negotiated Execution API version that honours the ``only_failed`` clear scope. +# Cadwyn migrates unknown request fields forward instead of rejecting them, so an older core +# would silently perform a whole-run clear. The failed-only emission is therefore gated here and +# raises explicitly when the negotiated core is older, rather than silently clearing the whole run. +MIN_VERSION_ONLY_FAILED_CLEAR = "2026-11-13" + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.api.datamodels._generated import API_VERSION +else: + API_VERSION = "" + if TYPE_CHECKING: from sqlalchemy.orm.session import Session @@ -139,6 +150,25 @@ class TriggerDagRunOperator(BaseOperator): DAG run conf is immutable and will not be reset on rerun of an existing DAG run. When reset_dag_run=False and dag run exists, DagRunAlreadyExists will be raised. When reset_dag_run=True and dag run exists, existing DAG run will be cleared to rerun. + :param auto_clear_failed_tasks: Opt-in, off by default. When set to ``True`` and the triggered + child run already exists in a terminal failed state, only the failed tasks (and their + downstream) of that run are cleared instead of resetting the whole run, so already-succeeded + upstream tasks are preserved and are not re-run. The clear happens on the next execution of + this operator (for example on a task retry): the failed child run is cleared and re-run at + that point, not asynchronously in the background -- so bound the number of attempts with + ``retries``. Because the previously-succeeded tasks are not re-executed and the failed ones + are, the triggered Dag's tasks should be idempotent; re-running a non-idempotent failed task + (or its downstream) may produce duplicate side effects. ``reset_dag_run`` takes precedence: + if both ``reset_dag_run`` and ``auto_clear_failed_tasks`` are set, ``reset_dag_run`` wins and + the whole run is cleared (ADR-017 precedence); at most one clear is ever performed. This is a + v1, synchronous-only feature: it applies when ``wait_for_completion=True`` and is not + supported with ``deferrable=True`` (setting both raises ``ValueError``). Note the brief + cosmetic window in which the child run may still be shown as ``failed`` until this operator + re-executes and triggers the clear. On Airflow 3.x the failed-only clear is delivered + server-side via the Execution API and therefore requires an Execution API version of at least + ``2026-11-13``; against an older core the operator refuses to fall back to a whole-run clear + and raises ``NotImplementedError`` rather than silently clearing more than requested. + (default: False) :param wait_for_completion: Whether or not wait for DAG run completion. (default: False) :param poke_interval: Poke interval to check DAG run status when wait_for_completion=True. (default: 60) @@ -190,6 +220,7 @@ def __init__( logical_date: str | datetime.datetime | None | ArgNotSet = NOTSET, run_after: str | datetime.datetime | None | ArgNotSet = NOTSET, reset_dag_run: bool = False, + auto_clear_failed_tasks: bool = False, wait_for_completion: bool = False, poke_interval: int = 60, allowed_states: list[str | DagRunState] | None = None, @@ -201,11 +232,18 @@ def __init__( openlineage_inject_parent_info: bool = True, **kwargs, ) -> None: + if not isinstance(auto_clear_failed_tasks, bool): + raise TypeError( + f"auto_clear_failed_tasks must be a bool, got {type(auto_clear_failed_tasks).__name__}" + ) + if auto_clear_failed_tasks and deferrable: + raise ValueError("auto_clear_failed_tasks is not supported with deferrable=True in this version.") super().__init__(**kwargs) self.trigger_dag_id = trigger_dag_id self.trigger_run_id = trigger_run_id self.conf = conf self.reset_dag_run = reset_dag_run + self.auto_clear_failed_tasks = auto_clear_failed_tasks self.wait_for_completion = wait_for_completion self.poke_interval = poke_interval if allowed_states: @@ -325,6 +363,22 @@ def _trigger_dag_af_3(self, context, run_id, parsed_logical_date, parsed_run_aft if parsed_run_after and "run_after" in parameters: kwargs_accepted["run_after"] = parsed_run_after + if self._failed_only_clear_requested(): + if API_VERSION < MIN_VERSION_ONLY_FAILED_CLEAR: + raise NotImplementedError( + "auto_clear_failed_tasks requires an Execution API of at least " + f"{MIN_VERSION_ONLY_FAILED_CLEAR}; the negotiated version is {API_VERSION}. " + "Refusing to fall back to a whole-run clear." + ) + if "only_failed_and_downstream" in parameters: + kwargs_accepted["only_failed_and_downstream"] = True + self.log.info( + "Requesting failed-only clear of existing run: dag_id=%s run_id=%s reason=%s", + self.trigger_dag_id, + run_id, + "auto_clear_failed_tasks set; only failed and downstream tasks will be cleared", + ) + if isinstance(context, Mapping): from airflow.utils import helpers @@ -344,6 +398,40 @@ def _trigger_dag_af_3(self, context, run_id, parsed_logical_date, parsed_run_aft raise DagRunTriggerException(**kwargs_accepted) + def _child_run_terminally_failed(self, dag_run) -> bool: + """Return True when the existing child run is in a terminal failed state.""" + return dag_run.state in self.failed_states + + def _failed_only_clear_requested(self) -> bool: + """ + Return True when the operator asks for a failed-only clear (ADR-017 precedence). + + ``reset_dag_run`` wins over ``auto_clear_failed_tasks`` when both are set, so the + failed-only intent is requested only when auto-clear is on and reset is off. Shared by + the AF2 action resolver and the AF3 signal emission so both paths decide precedence once. + """ + if self.reset_dag_run and self.auto_clear_failed_tasks: + self.log.warning( + "Both reset_dag_run and auto_clear_failed_tasks are set; reset_dag_run takes " + "precedence and the whole run will be cleared." + ) + return self.auto_clear_failed_tasks and not self.reset_dag_run + + def _resolve_already_exists_action(self, dag_run) -> str: + """ + Resolve how to handle an already-existing child run (ADR-017 precedence ladder). + + ``reset_dag_run`` wins over ``auto_clear_failed_tasks`` when both are set, so at most + one clear is ever performed. Returns ``"reset"`` (whole-run clear), ``"auto_clear"`` + (failed-only clear of a terminal-failed run), or ``"legacy"`` (existing skip/raise). + """ + if self.reset_dag_run: + self._failed_only_clear_requested() + return "reset" + if self._failed_only_clear_requested() and self._child_run_terminally_failed(dag_run): + return "auto_clear" + return "legacy" + def _trigger_dag_af_2(self, context, run_id, parsed_logical_date): try: unsupported_parameters = [] @@ -366,8 +454,9 @@ def _trigger_dag_af_2(self, context, run_id, parsed_logical_date): ) except DagRunAlreadyExists as e: - if self.reset_dag_run: - dag_run = e.dag_run + dag_run = e.dag_run + action = self._resolve_already_exists_action(dag_run) + if action == "reset": self.log.info("Clearing %s on %s", self.trigger_dag_id, dag_run.run_id) # Get target dag object and call clear() @@ -378,6 +467,17 @@ def _trigger_dag_af_2(self, context, run_id, parsed_logical_date): # Note: here execution fails on database isolation mode. Needs structural changes for AIP-72 dag = SerializedDagModel.get_dag(self.trigger_dag_id) dag.clear(start_date=dag_run.logical_date, end_date=dag_run.logical_date) + elif action == "auto_clear": + self.log.info( + "Auto-clearing failed tasks of existing run: dag_id=%s run_id=%s reason=%s", + self.trigger_dag_id, + dag_run.run_id, + "terminal-failed run cleared with only_failed=True", + ) + + # Note: here execution fails on database isolation mode. Needs structural changes for AIP-72 + dag = SerializedDagModel.get_dag(self.trigger_dag_id) + dag.clear(run_id=dag_run.run_id, only_failed=True) else: if self.skip_when_already_exists: raise AirflowSkipException( diff --git a/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py b/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py index 0d061fb75b57f..d8cec9ab56d1f 100644 --- a/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py +++ b/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py @@ -643,6 +643,352 @@ def test_trigger_dagrun_conf_openlineage_injection_valid_data( # Verify _get_openlineage_parent_info was called with ti mock_get_parent_info.assert_called_once_with(ti=mock_ti) + def test_auto_clear_failed_tasks_defaults_to_false(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + ) + + assert task.auto_clear_failed_tasks is False # @AC-FR001-01 + + def test_auto_clear_failed_tasks_accepts_true(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=True, + ) + + assert task.auto_clear_failed_tasks is True # @AC-FR001-01 + + def test_auto_clear_failed_tasks_non_bool_raises_type_error(self): + with pytest.raises(TypeError, match="auto_clear_failed_tasks must be a bool"): + TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks="yes", + ) # @AC-FR001-03 + + def test_auto_clear_failed_tasks_with_deferrable_raises_value_error(self): + with pytest.raises(ValueError, match="auto_clear_failed_tasks is not supported with deferrable"): + TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=True, + deferrable=True, + ) # @AC-FR008-01 + + @pytest.mark.parametrize( + ("state", "expected"), + [ + (DagRunState.FAILED, True), + (DagRunState.SUCCESS, False), + (DagRunState.RUNNING, False), + (DagRunState.QUEUED, False), + ], + ) + def test_child_run_terminally_failed_only_true_for_failed_state(self, state, expected): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + ) + dag_run = mock.MagicMock() + dag_run.state = state + + assert task._child_run_terminally_failed(dag_run) is expected # @AC-FR009-02 + + def test_child_run_terminally_failed_honours_custom_failed_states(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + failed_states=[DagRunState.QUEUED], + ) + dag_run = mock.MagicMock() + dag_run.state = DagRunState.QUEUED + + assert task._child_run_terminally_failed(dag_run) is True + + def test_reset_dag_run_wins_precedence_when_both_flags_set(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + reset_dag_run=True, + auto_clear_failed_tasks=True, + ) + dag_run = mock.MagicMock() + dag_run.state = DagRunState.FAILED + + assert task._resolve_already_exists_action(dag_run) == "reset" # @AC-FR007-01 + + def test_auto_clear_resolved_when_only_auto_clear_and_failed(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=True, + ) + dag_run = mock.MagicMock() + dag_run.state = DagRunState.FAILED + + assert task._resolve_already_exists_action(dag_run) == "auto_clear" # @AC-FR002-01 + + def test_legacy_action_when_auto_clear_but_run_not_terminally_failed(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=True, + ) + dag_run = mock.MagicMock() + dag_run.state = DagRunState.RUNNING + + assert task._resolve_already_exists_action(dag_run) == "legacy" # @AC-FR009-01 + + def test_legacy_action_when_no_flags_set(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + ) + dag_run = mock.MagicMock() + dag_run.state = DagRunState.FAILED + + assert task._resolve_already_exists_action(dag_run) == "legacy" + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="AF3 signal emission path") + @pytest.mark.parametrize( + ("reset_dag_run", "auto_clear_failed_tasks", "expected"), + [ + (False, False, False), + (True, False, False), + (False, True, True), + (True, True, False), + ], + ) + def test_af3_emits_only_failed_and_downstream_for_each_precedence_combination( + self, reset_dag_run, auto_clear_failed_tasks, expected + ): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + reset_dag_run=reset_dag_run, + auto_clear_failed_tasks=auto_clear_failed_tasks, + openlineage_inject_parent_info=False, + ) + + with mock.patch(f"{TRIGGER_OP_PATH}.API_VERSION", "2999-01-01"): + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) + + assert exc_info.value.only_failed_and_downstream is expected + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="AF3 signal emission path") + def test_af3_reset_dag_run_wins_over_auto_clear_when_both_set(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + reset_dag_run=True, + auto_clear_failed_tasks=True, + openlineage_inject_parent_info=False, + ) + + with mock.patch(f"{TRIGGER_OP_PATH}.API_VERSION", "2999-01-01"): + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) + + assert exc_info.value.reset_dag_run is True + assert exc_info.value.only_failed_and_downstream is False + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="AF3 signal emission path") + def test_af3_failed_only_emission_raises_when_core_below_capability_marker(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=True, + openlineage_inject_parent_info=False, + ) + + with mock.patch(f"{TRIGGER_OP_PATH}.API_VERSION", "2026-11-12"): + with pytest.raises(NotImplementedError, match="auto_clear_failed_tasks"): + task.execute(context={}) + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="AF3 signal emission path") + def test_af3_whole_run_emission_not_gated_on_capability_marker(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + reset_dag_run=True, + openlineage_inject_parent_info=False, + ) + + with mock.patch(f"{TRIGGER_OP_PATH}.API_VERSION", "2000-01-01"): + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) + + assert exc_info.value.reset_dag_run is True + assert exc_info.value.only_failed_and_downstream is False + + @pytest.mark.parametrize( + ("auto_clear_failed_tasks", "wait_for_completion"), + [ + pytest.param(True, False, id="defer-guard-only"), + pytest.param(True, True, id="defer-guard-with-wait"), + ], + ) + def test_auto_clear_with_deferrable_fails_fast_at_init( + self, auto_clear_failed_tasks, wait_for_completion + ): + with pytest.raises(ValueError, match="auto_clear_failed_tasks is not supported with deferrable"): + TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=auto_clear_failed_tasks, + wait_for_completion=wait_for_completion, + deferrable=True, + ) # @AC-FR008-01 + + def test_auto_clear_with_deferrable_never_defers(self): + # The guarded config must fail fast at construction and never reach a deferral, so no + # DagStateTrigger deferral / deferred DagRun / triggerer slot can result. + with mock.patch.object(TriggerDagRunOperator, "defer", autospec=True) as mock_defer: + with pytest.raises(ValueError, match="auto_clear_failed_tasks is not supported with deferrable"): + TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=True, + wait_for_completion=True, + deferrable=True, + ) + + mock_defer.assert_not_called() # @AC-FR008-02 + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="AF3 signal emission path") + def test_auto_clear_sync_mode_proceeds_normally(self): + # Control case: with deferrable=False the guard does not fire and the sync path runs. + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=True, + deferrable=False, + openlineage_inject_parent_info=False, + ) + + assert task.auto_clear_failed_tasks is True + assert task.deferrable is False + + with mock.patch.object(TriggerDagRunOperator, "defer", autospec=True) as mock_defer: + with mock.patch(f"{TRIGGER_OP_PATH}.API_VERSION", "2999-01-01"): + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) + + mock_defer.assert_not_called() + assert exc_info.value.only_failed_and_downstream is True # @AC-FR008-02 + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="AF3 signal emission path") + def test_off_path_emits_whole_run_signal_identically(self): + # With auto_clear off, the emitted signal must be byte-for-byte the legacy whole-run request. + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=False, + openlineage_inject_parent_info=False, + ) + + with mock.patch(f"{TRIGGER_OP_PATH}.API_VERSION", "2999-01-01"): + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) + + assert exc_info.value.only_failed_and_downstream is False # @AC-FR002-04 + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="AF3 signal emission path") + def test_off_path_does_not_gate_on_capability_marker(self): + # Off-path must never consult the capability marker, so an old core still triggers normally. + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=False, + openlineage_inject_parent_info=False, + ) + + with mock.patch(f"{TRIGGER_OP_PATH}.API_VERSION", "2000-01-01"): + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) + + assert exc_info.value.only_failed_and_downstream is False + + def test_resolver_yields_exactly_one_clear_across_precedence_matrix(self): + # Whole grid of (reset_dag_run, auto_clear_failed_tasks): at most one clear, reset wins. + expected = { + (False, False): "legacy", + (True, False): "reset", + (False, True): "auto_clear", + (True, True): "reset", + } + dag_run = mock.MagicMock() + dag_run.state = DagRunState.FAILED + for (reset_dag_run, auto_clear_failed_tasks), action in expected.items(): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + reset_dag_run=reset_dag_run, + auto_clear_failed_tasks=auto_clear_failed_tasks, + ) + assert task._resolve_already_exists_action(dag_run) == action # @AC-FR007-01 + + @pytest.mark.parametrize( + ("state", "expected_action"), + [ + (DagRunState.RUNNING, "legacy"), + (DagRunState.QUEUED, "legacy"), + (DagRunState.SUCCESS, "legacy"), + (DagRunState.FAILED, "auto_clear"), + ], + ) + def test_auto_clear_only_acts_on_terminal_failed_run(self, state, expected_action): + # Non-terminal / non-failed child never resolves to a clear; only a failed run does. + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=True, + ) + dag_run = mock.MagicMock() + dag_run.state = state + + assert task._resolve_already_exists_action(dag_run) == expected_action # @AC-FR009-01 + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="AF3 signal emission path") + def test_auto_clear_emits_audit_log_with_dag_and_run_id(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="custom_run_id", + auto_clear_failed_tasks=True, + openlineage_inject_parent_info=False, + ) + + with mock.patch.object(task.log, "info") as mock_info: + with mock.patch(f"{TRIGGER_OP_PATH}.API_VERSION", "2999-01-01"): + with pytest.raises(DagRunTriggerException): + task.execute(context={}) + + audit_calls = [c for c in mock_info.mock_calls if "failed-only clear" in str(c.args[0])] + assert len(audit_calls) == 1 + logged = audit_calls[0].args + assert TRIGGERED_DAG_ID in logged + assert "custom_run_id" in logged # @AC-FR006-01 + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="AF3 signal emission path") + def test_off_path_emits_no_auto_clear_audit_log(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + auto_clear_failed_tasks=False, + openlineage_inject_parent_info=False, + ) + + with mock.patch.object(task.log, "info") as mock_info: + with mock.patch(f"{TRIGGER_OP_PATH}.API_VERSION", "2999-01-01"): + with pytest.raises(DagRunTriggerException): + task.execute(context={}) + + assert not [c for c in mock_info.mock_calls if "failed-only clear" in str(c.args[0])] # @AC-FR006-02 + # TODO: To be removed once the provider drops support for Airflow 2 @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 2") @@ -940,6 +1286,108 @@ def test_trigger_dagrun_with_reset_dag_run_true( assert len(dag_runs) == expected_dagruns_count assert dag_runs[0].external_trigger + def test_auto_clear_failed_tasks_clears_only_failed_for_terminal_run(self, dag_maker): + logical_date = DEFAULT_DATE + run_id = "dummy_run_id" + with dag_maker( + TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True + ): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id=run_id, + logical_date=logical_date, + auto_clear_failed_tasks=True, + ) + dag_maker.sync_dagbag_to_db() + parse_and_sync_to_db(self.f_name) + dag_maker.create_dagrun() + dag_run = DagRun( + dag_id=TRIGGERED_DAG_ID, + execution_date=logical_date, + state=DagRunState.FAILED, + run_type="manual", + run_id=run_id, + ) + dag_maker.session.add(dag_run) + dag_maker.session.commit() + + with mock.patch("airflow.models.serialized_dag.SerializedDagModel.get_dag") as mock_get_dag: + task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) + + mock_get_dag.return_value.clear.assert_called_once_with( + run_id=run_id, only_failed=True + ) # @AC-FR002-01 + + dag_runs = dag_maker.session.scalars(select(DagRun).where(DagRun.dag_id == TRIGGERED_DAG_ID)).all() + assert len(dag_runs) == 1 + + def test_auto_clear_failed_tasks_noop_when_run_not_failed(self, dag_maker): + logical_date = DEFAULT_DATE + run_id = "dummy_run_id" + with dag_maker( + TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True + ): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id=run_id, + logical_date=logical_date, + auto_clear_failed_tasks=True, + skip_when_already_exists=True, + ) + dag_maker.sync_dagbag_to_db() + parse_and_sync_to_db(self.f_name) + dag_maker.create_dagrun() + dag_run = DagRun( + dag_id=TRIGGERED_DAG_ID, + execution_date=logical_date, + state=DagRunState.SUCCESS, + run_type="manual", + run_id=run_id, + ) + dag_maker.session.add(dag_run) + dag_maker.session.commit() + + with mock.patch("airflow.models.serialized_dag.SerializedDagModel.get_dag") as mock_get_dag: + task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) + + mock_get_dag.return_value.clear.assert_not_called() # @AC-FR009-01 + + def test_reset_dag_run_wins_over_auto_clear(self, dag_maker): + logical_date = DEFAULT_DATE + run_id = "dummy_run_id" + with dag_maker( + TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True + ): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id=run_id, + logical_date=logical_date, + reset_dag_run=True, + auto_clear_failed_tasks=True, + ) + dag_maker.sync_dagbag_to_db() + parse_and_sync_to_db(self.f_name) + dag_maker.create_dagrun() + dag_run = DagRun( + dag_id=TRIGGERED_DAG_ID, + execution_date=logical_date, + state=DagRunState.FAILED, + run_type="manual", + run_id=run_id, + ) + dag_maker.session.add(dag_run) + dag_maker.session.commit() + + with mock.patch("airflow.models.serialized_dag.SerializedDagModel.get_dag") as mock_get_dag: + task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) + + mock_get_dag.return_value.clear.assert_called_once_with( + start_date=dag_run.logical_date, end_date=dag_run.logical_date + ) # @AC-FR007-02 + def test_trigger_dagrun_with_wait_for_completion_true(self, dag_maker): """Test TriggerDagRunOperator with wait_for_completion.""" logical_date = DEFAULT_DATE @@ -1396,3 +1844,118 @@ def test_trigger_dagrun_conf_openlineage_injection_valid_data( assert dagrun.conf == injected_conf # Verify _get_openlineage_parent_info was called mock_get_parent_info.assert_called_once() + + def _seed_existing_run(self, dag_maker, run_id, logical_date, state): + dag_maker.sync_dagbag_to_db() + parse_and_sync_to_db(self.f_name) + dag_maker.create_dagrun() + dag_run = DagRun( + dag_id=TRIGGERED_DAG_ID, + execution_date=logical_date, + state=state, + run_type="manual", + run_id=run_id, + ) + dag_maker.session.add(dag_run) + dag_maker.session.commit() + return dag_run + + def test_off_path_never_clears_and_matches_legacy_raise(self, dag_maker): + # auto_clear_failed_tasks=False must behave exactly like today: no clear, DagRunAlreadyExists. + logical_date = DEFAULT_DATE + run_id = "dummy_run_id" + with dag_maker( + TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True + ): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id=run_id, + logical_date=logical_date, + auto_clear_failed_tasks=False, + ) + self._seed_existing_run(dag_maker, run_id, logical_date, DagRunState.FAILED) + + with mock.patch("airflow.models.serialized_dag.SerializedDagModel.get_dag") as mock_get_dag: + with pytest.raises(DagRunAlreadyExists): + task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) + + mock_get_dag.return_value.clear.assert_not_called() # @AC-FR002-04 + + def test_combined_mixed_run_clears_only_failed_and_downstream(self, dag_maker): + # Succeeded upstream is preserved (whole-run clear NOT used); only failed+downstream cleared. + logical_date = DEFAULT_DATE + run_id = "dummy_run_id" + with dag_maker( + TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True + ): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id=run_id, + logical_date=logical_date, + auto_clear_failed_tasks=True, + ) + self._seed_existing_run(dag_maker, run_id, logical_date, DagRunState.FAILED) + + with mock.patch("airflow.models.serialized_dag.SerializedDagModel.get_dag") as mock_get_dag: + task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) + + mock_clear = mock_get_dag.return_value.clear + mock_clear.assert_called_once_with(run_id=run_id, only_failed=True) # @AC-FR003-02 + # A whole-run clear (start_date/end_date) would wipe succeeded upstream — it must not be used. + _, kwargs = mock_clear.call_args + assert "start_date" not in kwargs and "end_date" not in kwargs + + dag_runs = dag_maker.session.scalars(select(DagRun).where(DagRun.dag_id == TRIGGERED_DAG_ID)).all() + assert len(dag_runs) == 1 # @AC-FR004-01 + + def test_repeated_failure_is_bounded_and_never_falsely_succeeds(self, dag_maker): + # Re-running against a still-failed run clears once per execute; it never loops or flips to success. + logical_date = DEFAULT_DATE + run_id = "dummy_run_id" + with dag_maker( + TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True + ): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id=run_id, + logical_date=logical_date, + auto_clear_failed_tasks=True, + ) + self._seed_existing_run(dag_maker, run_id, logical_date, DagRunState.FAILED) + + with mock.patch("airflow.models.serialized_dag.SerializedDagModel.get_dag") as mock_get_dag: + task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) + task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) + + # Each execute performs at most one clear; no re-clear loop within a single execute. + assert mock_get_dag.return_value.clear.call_count == 2 # @AC-FR005-01 + + dag_runs = dag_maker.session.scalars(select(DagRun).where(DagRun.dag_id == TRIGGERED_DAG_ID)).all() + assert len(dag_runs) == 1 # @AC-FR005-02 + + def test_auto_clear_writes_audit_log_row(self, dag_maker): + logical_date = DEFAULT_DATE + run_id = "dummy_run_id" + with dag_maker( + TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True + ): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id=run_id, + logical_date=logical_date, + auto_clear_failed_tasks=True, + ) + self._seed_existing_run(dag_maker, run_id, logical_date, DagRunState.FAILED) + + with mock.patch("airflow.models.serialized_dag.SerializedDagModel.get_dag"): + with mock.patch.object(task.log, "info") as mock_info: + task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) + + audit_calls = [c for c in mock_info.mock_calls if "Auto-clearing failed tasks" in str(c.args[0])] + assert len(audit_calls) == 1 + logged = audit_calls[0].args + assert TRIGGERED_DAG_ID in logged and run_id in logged # @AC-FR006-01 diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 42afaace779d8..aa840fd64ca90 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -50,6 +50,7 @@ AssetResponse, AssetStateStorePutBody, AssetStateStoreResponse, + ClearDagRunPayload, ConnectionResponse, ConnectionTestConnectionResponse, ConnectionTestResultBody, @@ -975,6 +976,8 @@ def trigger( run_after: datetime | None = None, reset_dag_run: bool = False, note: str | None = None, + *, + only_failed: bool = False, ) -> OKResponse | ErrorResponse: """Trigger a Dag run via the API server.""" body = TriggerDAGRunPayload( @@ -991,9 +994,9 @@ def trigger( ) except ServerResponseError as e: if e.response.status_code == HTTPStatus.CONFLICT: - if reset_dag_run: - log.info("Dag Run already exists; Resetting Dag Run.", dag_id=dag_id, run_id=run_id) - return self.clear(run_id=run_id, dag_id=dag_id) + if reset_dag_run or only_failed: + log.info("Dag Run already exists; clearing to rerun.", dag_id=dag_id, run_id=run_id) + return self.clear(run_id=run_id, dag_id=dag_id, only_failed=only_failed) log.info("Dag Run already exists!", detail=e.detail, dag_id=dag_id, run_id=run_id) return ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) @@ -1001,9 +1004,12 @@ def trigger( return OKResponse(ok=True) - def clear(self, dag_id: str, run_id: str) -> OKResponse: + def clear(self, dag_id: str, run_id: str, *, only_failed: bool = False) -> OKResponse: """Clear a Dag run via the API server.""" - self.client.post(f"dag-runs/{dag_id}/{run_id}/clear") + body = ClearDagRunPayload(only_failed=only_failed) + self.client.post( + f"dag-runs/{dag_id}/{run_id}/clear", content=body.model_dump_json(exclude_defaults=True) + ) # TODO: Error handling return OKResponse(ok=True) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 02d51e15b57e5..73dfc841ab816 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -597,6 +597,17 @@ class TriggerDAGRunPayload(BaseModel): note: Annotated[str | None, Field(title="Note")] = None +class ClearDagRunPayload(BaseModel): + """ + Schema for Clear DAG Run API request. + """ + + model_config = ConfigDict( + extra="forbid", + ) + only_failed: Annotated[bool | None, Field(title="Only Failed")] = False + + class UpdateHITLDetailPayload(BaseModel): """ Schema for writing the response part of a Human-in-the-loop detail for a specific task instance. diff --git a/task-sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py index 78a7c491c999b..e490eb12906dd 100644 --- a/task-sdk/src/airflow/sdk/exceptions.py +++ b/task-sdk/src/airflow/sdk/exceptions.py @@ -303,6 +303,7 @@ def __init__( logical_date=None, run_after=None, reset_dag_run: bool, + only_failed_and_downstream: bool = False, skip_when_already_exists: bool, wait_for_completion: bool, allowed_states: list[str], @@ -318,6 +319,7 @@ def __init__( self.logical_date = logical_date self.run_after = run_after self.reset_dag_run = reset_dag_run + self.only_failed_and_downstream = only_failed_and_downstream self.skip_when_already_exists = skip_when_already_exists self.wait_for_completion = wait_for_completion self.allowed_states = allowed_states diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 7c957027ba792..70a26f30767c7 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -1139,6 +1139,7 @@ class SetExecutionTimeout(BaseModel): class TriggerDagRun(TriggerDAGRunPayload): dag_id: str run_id: Annotated[str, Field(title="Dag Run Id")] + only_failed_and_downstream: bool = False type: Literal["TriggerDagRun"] = "TriggerDagRun" diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 77888c5adfdea..e4d283adcc15a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1,6 +1,6 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-10-30", + "api_version": "2026-11-13", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { "AssetAliasReferenceAssetEventDagRun": { diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 92295bd4f8ddb..d8ed1d1828096 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -2206,7 +2206,14 @@ def _handle_trigger_dag_run( self, msg: TriggerDagRun, log: FilteringBoundLogger, req_id: int ) -> RequestResult: resp = self.client.dag_runs.trigger( - msg.dag_id, msg.run_id, msg.conf, msg.logical_date, msg.run_after, msg.reset_dag_run, msg.note + msg.dag_id, + msg.run_id, + msg.conf, + msg.logical_date, + msg.run_after, + msg.reset_dag_run, + msg.note, + only_failed=msg.only_failed_and_downstream, ) return resp, {} diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 88fb19293d0c1..81449eea5a0a8 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -2028,6 +2028,7 @@ def _handle_trigger_dag_run( run_after=drte.run_after, conf=drte.conf, reset_dag_run=drte.reset_dag_run, + only_failed_and_downstream=drte.only_failed_and_downstream, note=drte.note, ), ) diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 62fde3373e007..5b2eeffebc9f2 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -1669,6 +1669,8 @@ def test_clear(self): def handle_request(request: httpx2.Request) -> httpx2.Response: if request.url.path == "/dag-runs/test_clear/test_run_id/clear": + # Default whole-run clear must not carry an only_failed body. + assert "only_failed" not in json.loads(request.read() or "{}") return httpx2.Response(status_code=204) return httpx2.Response(status_code=422) @@ -1677,6 +1679,86 @@ def handle_request(request: httpx2.Request) -> httpx2.Response: assert result == OKResponse(ok=True) + def test_clear_only_failed(self): + """Test that clear(only_failed=True) scopes the server clear to failed+downstream.""" + + def handle_request(request: httpx2.Request) -> httpx2.Response: + if request.url.path == "/dag-runs/test_clear/test_run_id/clear": + assert json.loads(request.read())["only_failed"] is True + return httpx2.Response(status_code=204) + return httpx2.Response(status_code=422) + + client = make_client(transport=httpx2.MockTransport(handle_request)) + result = client.dag_runs.clear(dag_id="test_clear", run_id="test_run_id", only_failed=True) + + assert result == OKResponse(ok=True) + + def test_trigger_conflict_reset_dag_run_only_failed(self): + """Test that the trigger->clear fallback forwards only_failed to the clear body.""" + + def handle_request(request: httpx2.Request) -> httpx2.Response: + if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id": + return httpx2.Response( + status_code=409, + json={ + "detail": { + "reason": "already_exists", + "message": "A Dag Run already exists for Dag test_trigger_conflict with run id test_run_id", + } + }, + ) + if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id/clear": + assert json.loads(request.read())["only_failed"] is True + return httpx2.Response(status_code=204) + return httpx2.Response(status_code=422) + + client = make_client(transport=httpx2.MockTransport(handle_request)) + result = client.dag_runs.trigger( + dag_id="test_trigger_conflict_reset", + run_id="test_run_id", + reset_dag_run=True, + only_failed=True, + ) + + assert result == OKResponse(ok=True) + + def test_trigger_conflict_only_failed_without_reset(self): + """Regression: only_failed=True must reach the clear even when reset_dag_run=False. + + This is the auto_clear_failed_tasks path (reset_dag_run stays False). A conflict + must be resolved by a failed-only clear rather than returning + ``DAGRUN_ALREADY_EXISTS`` -- gating the clear on ``reset_dag_run`` alone makes the + whole feature a no-op. + """ + clear_body: dict = {} + + def handle_request(request: httpx2.Request) -> httpx2.Response: + if request.url.path == "/dag-runs/test_trigger_conflict_only_failed/test_run_id": + return httpx2.Response( + status_code=409, + json={ + "detail": { + "reason": "already_exists", + "message": "A Dag Run already exists for Dag test_trigger_conflict_only_failed", + } + }, + ) + if request.url.path == "/dag-runs/test_trigger_conflict_only_failed/test_run_id/clear": + clear_body.update(json.loads(request.read())) + return httpx2.Response(status_code=204) + return httpx2.Response(status_code=422) + + client = make_client(transport=httpx2.MockTransport(handle_request)) + result = client.dag_runs.trigger( + dag_id="test_trigger_conflict_only_failed", + run_id="test_run_id", + reset_dag_run=False, + only_failed=True, + ) + + assert result == OKResponse(ok=True) + assert clear_body == {"only_failed": True} + def test_get_state(self): """Test that the client can get the state of a dag run""" diff --git a/task-sdk/tests/task_sdk/execution_time/test_comms.py b/task-sdk/tests/task_sdk/execution_time/test_comms.py index 88f8f73363e62..ac0a2029573cb 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_comms.py +++ b/task-sdk/tests/task_sdk/execution_time/test_comms.py @@ -33,6 +33,7 @@ GetVariable, MaskSecret, StartupDetails, + TriggerDagRun, VariableResult, _RequestFrame, _ResponseFrame, @@ -62,6 +63,18 @@ def test_mask_secret_with_objects(self, object_to_mask): mask_secret_object = MaskSecret(value=object_to_mask, name="test_secret") assert mask_secret_object.value == object_to_mask + def test_trigger_dag_run_only_failed_and_downstream_defaults_false(self): + msg = TriggerDagRun(dag_id="d", run_id="r") + assert msg.only_failed_and_downstream is False # @AC-FR002-01 + + @pytest.mark.parametrize("only_failed_and_downstream", [True, False]) + def test_trigger_dag_run_only_failed_and_downstream_round_trips(self, only_failed_and_downstream): + msg = TriggerDagRun( + dag_id="d", run_id="r", only_failed_and_downstream=only_failed_and_downstream + ) + restored = TriggerDagRun.model_validate_json(msg.model_dump_json()) + assert restored.only_failed_and_downstream is only_failed_and_downstream + class TestCommsDecoder: """Test the communication between the subprocess and the "supervisor".""" diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 433d6d11f3685..15b025a9a9f1f 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -3295,6 +3295,7 @@ class RequestTestCase: True, None, ), + kwargs={"only_failed": False}, response=OKResponse(ok=True), ), test_id="dag_run_trigger_with_run_after", @@ -3320,6 +3321,7 @@ class RequestTestCase: True, "Test Note", ), + kwargs={"only_failed": False}, response=OKResponse(ok=True), ), test_id="dag_run_trigger", @@ -3336,16 +3338,36 @@ class RequestTestCase: client_mock=ClientMock( method_path="dag_runs.trigger", args=("test_dag", "test_run", {"key": "value"}, timezone.datetime(2025, 1, 1), None, True, None), + kwargs={"only_failed": False}, response=OKResponse(ok=True), ), test_id="dag_run_trigger", ), + RequestTestCase( + message=TriggerDagRun( + dag_id="test_dag", + run_id="test_run", + conf={"key": "value"}, + logical_date=timezone.datetime(2025, 1, 1), + reset_dag_run=True, + only_failed_and_downstream=True, + ), + expected_body={"ok": True, "type": "OKResponse"}, + client_mock=ClientMock( + method_path="dag_runs.trigger", + args=("test_dag", "test_run", {"key": "value"}, timezone.datetime(2025, 1, 1), None, True, None), + kwargs={"only_failed": True}, + response=OKResponse(ok=True), + ), + test_id="dag_run_trigger_only_failed_and_downstream", + ), RequestTestCase( message=TriggerDagRun(dag_id="test_dag", run_id="test_run"), expected_body={"error": "DAGRUN_ALREADY_EXISTS", "detail": None, "type": "ErrorResponse"}, client_mock=ClientMock( method_path="dag_runs.trigger", args=("test_dag", "test_run", None, None, None, False, None), + kwargs={"only_failed": False}, response=ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS), ), test_id="dag_run_trigger_already_exists", diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index cbce1bb32d105..8f1762e6c677c 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -104,6 +104,7 @@ AirflowSkipException, AirflowTaskTerminated, AirflowTaskTimeout, + DagRunTriggerException, DownstreamTasksSkipped, ErrorType, TaskDeferred, @@ -5946,6 +5947,47 @@ def test_handle_trigger_dag_run(self, create_runtime_ti, mock_supervisor_comms): ] mock_supervisor_comms.assert_has_calls(expected_calls) + def test_dag_run_trigger_exception_only_failed_and_downstream_defaults_false(self): + exc = DagRunTriggerException( + trigger_dag_id="d", + dag_run_id="r", + conf=None, + reset_dag_run=False, + skip_when_already_exists=False, + wait_for_completion=False, + allowed_states=[], + failed_states=[], + poke_interval=1, + deferrable=False, + ) + assert exc.only_failed_and_downstream is False # @AC-FR002-01 + + @pytest.mark.parametrize("only_failed_and_downstream", [True, False]) + def test_handle_trigger_dag_run_propagates_only_failed_and_downstream( + self, only_failed_and_downstream, mock_supervisor_comms + ): + drte = DagRunTriggerException( + trigger_dag_id="test_dag", + dag_run_id="test_run_id", + conf=None, + reset_dag_run=False, + skip_when_already_exists=False, + wait_for_completion=False, + allowed_states=[], + failed_states=[], + poke_interval=1, + deferrable=False, + only_failed_and_downstream=only_failed_and_downstream, + ) + ti = mock.MagicMock() + ti.rendered_map_index = None + + task_runner._handle_trigger_dag_run(drte, mock.MagicMock(), ti, mock.MagicMock()) + + sent = mock_supervisor_comms.send.call_args_list[0].args[0] + assert isinstance(sent, TriggerDagRun) + assert sent.only_failed_and_downstream is only_failed_and_downstream + @pytest.mark.parametrize( ("skip_when_already_exists", "expected_state"), [ From 5e3865d319f7052ebad2077d9344500b6b51de6d Mon Sep 17 00:00:00 2001 From: Eugene Olkhovski <65240525+eugeneo17@users.noreply.github.com> Date: Fri, 25 Sep 2026 14:01:35 -0400 Subject: [PATCH 2/2] Add newsfragment for auto-clear of failed child tasks (#71730) Document the new Execution API version 2026-11-13 and the opt-in auto_clear_failed_tasks behavior on TriggerDagRunOperator. Co-Authored-By: Claude Opus 4.8 --- airflow-core/newsfragments/71730.feature.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 airflow-core/newsfragments/71730.feature.rst diff --git a/airflow-core/newsfragments/71730.feature.rst b/airflow-core/newsfragments/71730.feature.rst new file mode 100644 index 0000000000000..dad80764e9c60 --- /dev/null +++ b/airflow-core/newsfragments/71730.feature.rst @@ -0,0 +1 @@ +Add a new Execution API version (``2026-11-13``) that lets a running task clear only the failed tasks of an existing Dag run (and their downstream) instead of resetting the whole run. This backs the opt-in ``auto_clear_failed_tasks`` flag on ``TriggerDagRunOperator``: when a triggered child Dag run has already failed, re-running the parent task (for example on a retry) clears only the failed tasks so the child resumes from its point of failure while preserving already-succeeded tasks. ``reset_dag_run`` still takes precedence when both are set, and against an older core the operator raises ``NotImplementedError`` rather than silently falling back to a whole-run clear.