diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py index c8e7ab2378dd1..594423a0e0246 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py @@ -162,6 +162,11 @@ class ClearTaskInstancesBody(StrictBaseModel): include_downstream: bool = False include_future: bool = False include_past: bool = False + include_recursive: bool = Field( + default=False, + description="If set to true, clearing tasks that are ExternalTaskMarker instances " + "will also recursively clear the referenced tasks in external DAGs.", + ) run_on_latest_version: bool = Field( default=False, description="(Experimental) Run on the latest bundle version of the dag after " diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 0f98b70d39bfa..de7288513a214 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -9919,6 +9919,13 @@ components: type: boolean title: Include Past default: false + include_recursive: + type: boolean + title: Include Recursive + description: If set to true, clearing tasks that are ExternalTaskMarker + instances will also recursively clear the referenced tasks in external + DAGs. + default: false run_on_latest_version: type: boolean title: Run On Latest Version diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py index 1af56b565f309..34102e664b937 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py @@ -90,7 +90,11 @@ from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import AirflowClearRunningTaskException, TaskNotFound from airflow.models import Base, DagRun -from airflow.models.taskinstance import TaskInstance as TI, clear_task_instances +from airflow.models.taskinstance import ( + TaskInstance as TI, + clear_external_task_marker_dependencies, + clear_task_instances, +) from airflow.models.taskinstancehistory import TaskInstanceHistory as TIH from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS @@ -836,6 +840,16 @@ def _collect_relatives(run_id: str, direction: Literal["upstream", "downstream"] except AirflowClearRunningTaskException as e: raise HTTPException(status.HTTP_409_CONFLICT, str(e)) from e + # Recursively clear tasks in external DAGs referenced by ExternalTaskMarker + if body.include_recursive: + external_tis = clear_external_task_marker_dependencies( + tis=list(task_instances), + dag_bag=dag_bag, + session=session, + dag_run_state=DagRunState.QUEUED if reset_dag_runs else False, + ) + task_instances = list(task_instances) + external_tis + if body.note is not None: _patch_task_instance_note( task_instance_body=body, diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 6fb2b38696cf1..fe81de3f883f4 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -118,6 +118,7 @@ from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile from airflow.models.dag import DagModel + from airflow.models.dagbag import DBDagBag from airflow.models.dagrun import DagRun from airflow.serialization.definitions.dag import SerializedDAG from airflow.serialization.definitions.mappedoperator import Operator @@ -434,6 +435,143 @@ def clear_task_instances( session.flush() +def clear_external_task_marker_dependencies( + tis: list[TaskInstance], + dag_bag: DBDagBag, + session: Session, + dag_run_state: DagRunState | Literal[False] = DagRunState.QUEUED, + visited: set[tuple[str, str, str]] | None = None, + recursion_depth: int = 10, +) -> list[TaskInstance]: + """ + Recursively clear tasks referenced by ExternalTaskMarker instances. + + When clearing tasks, if any of the cleared task instances are backed by an + ExternalTaskMarker operator, this function will load the referenced external + DAG and clear the corresponding task(s) in that DAG. This process repeats + recursively up to ``recursion_depth`` levels. + + :param tis: task instances that were just cleared + :param dag_bag: DBDagBag used to load external DAGs + :param session: SQLAlchemy session + :param dag_run_state: state to set affected DagRuns to + :param visited: set of (dag_id, task_id, logical_date) already visited to prevent cycles + :param recursion_depth: maximum remaining recursion levels + :return: list of additionally cleared task instances from external DAGs + + :meta private: + """ + if recursion_depth <= 0: + raise RuntimeError( + f"Maximum recursion depth {recursion_depth} reached for ExternalTaskMarker. " + "This is likely caused by circular dependencies between DAGs." + ) + + if visited is None: + visited = set() + + all_external_tis: list[TaskInstance] = [] + + for ti in tis: + # Load the serialized DAG for this task instance to check its type + dag = dag_bag.get_latest_version_of_dag(ti.dag_id, session=session) + if not dag: + continue + + if not dag.has_task(ti.task_id): + continue + + task = dag.get_task(ti.task_id) + + # Check if this task is an ExternalTaskMarker by its task_type + if task.task_type != "ExternalTaskMarker": + continue + + external_dag_id = getattr(task, "external_dag_id", None) + external_task_id = getattr(task, "external_task_id", None) + marker_recursion_depth = getattr(task, "recursion_depth", 10) + + if not external_dag_id or not external_task_id: + continue + + # Use the marker's own recursion_depth as the limit, but respect the remaining depth + effective_depth = min(marker_recursion_depth, recursion_depth) + + # The logical_date for the external task — use the TI's logical_date + logical_date_key = str(ti.logical_date) if ti.logical_date else "" + visit_key = (external_dag_id, external_task_id, logical_date_key) + + if visit_key in visited: + continue + visited.add(visit_key) + + # Load the external DAG + external_dag = dag_bag.get_latest_version_of_dag(external_dag_id, session=session) + if not external_dag: + log.warning( + "ExternalTaskMarker on %s.%s references DAG '%s' which does not exist. Skipping.", + ti.dag_id, + ti.task_id, + external_dag_id, + ) + continue + + if not external_dag.has_task(external_task_id): + log.warning( + "ExternalTaskMarker on %s.%s references task '%s' in DAG '%s' which does not exist. Skipping.", + ti.dag_id, + ti.task_id, + external_task_id, + external_dag_id, + ) + continue + + # Find matching task instances in the external DAG + # Include the target task and all its downstream tasks + external_partial = external_dag.partial_subset( + task_ids=[external_task_id], + include_downstream=True, + ) + + external_tis = list( + external_dag._get_task_instances( + task_ids=list(external_partial.task_dict.keys()), + start_date=ti.logical_date, + end_date=ti.logical_date, + run_id=None, + state=(), + exclude_task_ids=None, + exclude_run_ids=None, + session=session, + ) + ) + + if not external_tis: + continue + + # Clear the external task instances + clear_task_instances( + external_tis, + session, + dag_run_state=dag_run_state, + ) + all_external_tis.extend(external_tis) + + # Recurse into the newly cleared tasks to handle transitive ExternalTaskMarker chains + if effective_depth > 1: + transitive_tis = clear_external_task_marker_dependencies( + tis=external_tis, + dag_bag=dag_bag, + session=session, + dag_run_state=dag_run_state, + visited=visited, + recursion_depth=effective_depth - 1, + ) + all_external_tis.extend(transitive_tis) + + return all_external_tis + + def _creator_note(val): """Creator for the ``note`` association proxy.""" if isinstance(val, str): 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 3f861d006219e..935c3279a247d 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 @@ -1373,6 +1373,12 @@ export const $ClearTaskInstancesBody = { title: 'Include Past', default: false }, + include_recursive: { + type: 'boolean', + title: 'Include Recursive', + description: 'If set to true, clearing tasks that are ExternalTaskMarker instances will also recursively clear the referenced tasks in external DAGs.', + default: false + }, run_on_latest_version: { type: 'boolean', title: 'Run On Latest Version', 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 adf7d7a0cc829..ca060cda13d25 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 @@ -422,6 +422,10 @@ export type ClearTaskInstancesBody = { include_downstream?: boolean; include_future?: boolean; include_past?: boolean; + /** + * If set to true, clearing tasks that are ExternalTaskMarker instances will also recursively clear the referenced tasks in external DAGs. + */ + include_recursive?: boolean; /** * (Experimental) Run on the latest bundle version of the dag after clearing the task instances. */ diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index d874b6163a039..6dab0c6be4395 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -3655,3 +3655,146 @@ def test_clear_task_instances_resets_context_carrier(dag_maker, session): assert ti.context_carrier["traceparent"] != original_ti_traceparent assert dag_run.context_carrier["traceparent"] != original_dr_traceparent + + +class TestClearExternalTaskMarkerDependencies: + """Tests for clear_external_task_marker_dependencies.""" + + def test_clears_external_dag_tasks(self, dag_maker, session): + """Test that clearing an ExternalTaskMarker clears the referenced task in the external DAG.""" + from airflow.models.dagbag import DBDagBag + from airflow.models.taskinstance import clear_external_task_marker_dependencies + from airflow.providers.standard.sensors.external_task import ExternalTaskMarker + + # Create the parent DAG with an ExternalTaskMarker + with dag_maker("parent_dag", session=session): + EmptyOperator(task_id="upstream_task") + ExternalTaskMarker( + task_id="marker_task", + external_dag_id="child_dag", + external_task_id="child_sensor", + ) + parent_dr = dag_maker.create_dagrun() + + # Create the child DAG + with dag_maker("child_dag", session=session): + EmptyOperator(task_id="child_sensor") + EmptyOperator(task_id="child_downstream") + child_dr = dag_maker.create_dagrun() + + # Set all TIs to success + for ti in parent_dr.task_instances: + ti.state = TaskInstanceState.SUCCESS + for ti in child_dr.task_instances: + ti.state = TaskInstanceState.SUCCESS + session.flush() + + # Get the marker TI + marker_ti = parent_dr.get_task_instance("marker_task", session=session) + + dag_bag = DBDagBag(load_op_links=False) + + # Call the function + external_tis = clear_external_task_marker_dependencies( + tis=[marker_ti], + dag_bag=dag_bag, + session=session, + ) + + # Verify external tasks were cleared + assert len(external_tis) > 0 + cleared_keys = {(ti.dag_id, ti.task_id) for ti in external_tis} + assert ("child_dag", "child_sensor") in cleared_keys + + def test_skips_non_marker_tasks(self, dag_maker, session): + """Test that non-ExternalTaskMarker tasks are ignored.""" + from airflow.models.dagbag import DBDagBag + from airflow.models.taskinstance import clear_external_task_marker_dependencies + + with dag_maker("regular_dag", session=session): + EmptyOperator(task_id="regular_task") + dr = dag_maker.create_dagrun() + + ti = dr.get_task_instance("regular_task", session=session) + ti.state = TaskInstanceState.SUCCESS + session.flush() + + dag_bag = DBDagBag(load_op_links=False) + + external_tis = clear_external_task_marker_dependencies( + tis=[ti], + dag_bag=dag_bag, + session=session, + ) + + assert external_tis == [] + + def test_handles_missing_external_dag(self, dag_maker, session): + """Test that a missing external DAG is handled gracefully.""" + from airflow.models.dagbag import DBDagBag + from airflow.models.taskinstance import clear_external_task_marker_dependencies + from airflow.providers.standard.sensors.external_task import ExternalTaskMarker + + with dag_maker("parent_dag_missing", session=session): + ExternalTaskMarker( + task_id="marker_task", + external_dag_id="nonexistent_dag", + external_task_id="nonexistent_task", + ) + dr = dag_maker.create_dagrun() + + ti = dr.get_task_instance("marker_task", session=session) + ti.state = TaskInstanceState.SUCCESS + session.flush() + + dag_bag = DBDagBag(load_op_links=False) + + # Should not raise, just skip + external_tis = clear_external_task_marker_dependencies( + tis=[ti], + dag_bag=dag_bag, + session=session, + ) + assert external_tis == [] + + def test_respects_recursion_depth(self, session): + """Test that recursion_depth limit raises RuntimeError.""" + from airflow.models.taskinstance import clear_external_task_marker_dependencies + + with pytest.raises(RuntimeError, match="Maximum recursion depth"): + clear_external_task_marker_dependencies( + tis=[], + dag_bag=mock.MagicMock(), + session=session, + recursion_depth=0, + ) + + def test_prevents_cycles(self, dag_maker, session): + """Test that visited set prevents infinite cycles.""" + from airflow.models.dagbag import DBDagBag + from airflow.models.taskinstance import clear_external_task_marker_dependencies + from airflow.providers.standard.sensors.external_task import ExternalTaskMarker + + with dag_maker("cyclic_dag", session=session): + ExternalTaskMarker( + task_id="marker_back", + external_dag_id="cyclic_dag", + external_task_id="marker_back", + ) + dr = dag_maker.create_dagrun() + + ti = dr.get_task_instance("marker_back", session=session) + ti.state = TaskInstanceState.SUCCESS + session.flush() + + dag_bag = DBDagBag(load_op_links=False) + + # Should not loop forever — visited set breaks the cycle + external_tis = clear_external_task_marker_dependencies( + tis=[ti], + dag_bag=dag_bag, + session=session, + ) + # The self-referencing marker should be cleared once + cleared_keys = {(t.dag_id, t.task_id) for t in external_tis} + assert ("cyclic_dag", "marker_back") in cleared_keys