Lee-W commented on code in PR #65314:
URL: https://github.com/apache/airflow/pull/65314#discussion_r3635978723
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -925,30 +930,69 @@ def _collect_relatives(run_id: str, direction:
Literal["upstream", "downstream"]
*((t, m) for t, m in mapped_tasks_tuples if t not in
normal_task_ids),
]
+ # Follow ExternalTaskMarker connections when explicitly requested via
include_downstream_dags, or
+ # automatically whenever downstream clearing is selected (restoring
Airflow 2 behavior)
+ include_dependent_dags = body.include_downstream_dags or downstream
+
task_instances: Sequence[TI]
- if dag_run_id is not None and not (past or future):
- # Use run_id-based clearing when we have a specific dag_run_id and not
using past/future
- task_instances = dag.clear(
- dry_run=True,
- task_ids=task_markers_to_clear,
- run_id=dag_run_id,
- session=session,
- run_on_latest_version=resolved_run_on_latest,
- only_failed=body.only_failed,
- only_running=body.only_running,
- )
- else:
- # Use date-based clearing when no dag_run_id or when past/future is
specified
- task_instances = dag.clear(
- dry_run=True,
- task_ids=task_markers_to_clear,
- start_date=body.start_date,
- end_date=body.end_date,
- session=session,
- run_on_latest_version=resolved_run_on_latest,
- only_failed=body.only_failed,
- only_running=body.only_running,
- )
+ try:
+ if dag_run_id is not None and not (past or future):
+ # Use run_id-based clearing when we have a specific dag_run_id and
not using past/future
+ task_instances = dag.clear(
+ dry_run=True,
+ task_ids=task_markers_to_clear,
+ run_id=dag_run_id,
+ session=session,
+ run_on_latest_version=resolved_run_on_latest,
+ only_failed=body.only_failed,
+ only_running=body.only_running,
+ include_dependent_dags=include_dependent_dags,
+ dag_bag=dag_bag,
+ )
+ else:
+ # Use date-based clearing when no dag_run_id or when past/future
is specified
+ task_instances = dag.clear(
+ dry_run=True,
+ task_ids=task_markers_to_clear,
+ start_date=body.start_date,
+ end_date=body.end_date,
+ session=session,
+ run_on_latest_version=resolved_run_on_latest,
+ only_failed=body.only_failed,
+ only_running=body.only_running,
+ include_dependent_dags=include_dependent_dags,
+ dag_bag=dag_bag,
+ )
+
+ except MaxRecursionDepthError as e:
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, str(e)) from e
+ except ParserError as e:
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, f"Invalid
logical_date: {e}") from e
+ except DagNotFound as e:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, str(e)) from e
+
+ if include_dependent_dags:
+ # Ensure proper access to downstream dags/tasks with dag.clear and
include_dependent_dags
Review Comment:
```suggestion
# Ensure proper access to downstream Dags/tasks with dag.clear and
include_dependent_dags
```
##########
airflow-core/tests/unit/serialization/definitions/test_dag.py:
##########
@@ -0,0 +1,387 @@
+#
+# 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 pendulum
+import pytest
+
+from airflow.exceptions import AirflowException
+from airflow.models.dagbag import DBDagBag
+from airflow.models.renderedtifields import RenderedTaskInstanceFields
+from airflow.providers.standard.sensors.external_task import
ExternalTaskMarker, ExternalTaskSensor
+
+from tests_common.test_utils import db
+from tests_common.test_utils.db import clear_rendered_ti_fields
+
+pytestmark = pytest.mark.db_test
+
+EXTERNAL_LOGICAL_DATE = pendulum.datetime(2024, 1, 1, tz="UTC")
+
+
[email protected](autouse=True)
+def reset_db():
+ db.clear_db_dags()
+ db.clear_db_runs()
+ db.clear_db_serialized_dags()
+ clear_rendered_ti_fields()
+
+
+def test_clear_does_not_follow_external_marker_by_default(dag_maker, session):
+ """Without include_dependent_dags, ExternalTaskMarker links are not
followed."""
+ with dag_maker("parent_dag", session=session, schedule=None):
+ ExternalTaskMarker(
+ task_id="trigger_child",
+ external_dag_id="child_dag",
+ external_task_id="wait_for_parent",
+ recursion_depth=3,
+ )
+
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ serialized_parent = dag_maker.serialized_dag
+
+ with dag_maker("child_dag", session=session, schedule=None):
+ ExternalTaskSensor(
+ task_id="wait_for_parent",
+ external_dag_id="parent_dag",
+ external_task_id="trigger_child",
+ poke_interval=5,
+ )
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ session.flush()
+
+ result = serialized_parent.clear(dry_run=True, only_failed=False,
session=session)
+
+ dag_ids = {ti.dag_id for ti in result}
+ assert "parent_dag" in dag_ids
+ assert "child_dag" not in dag_ids
+
+
+def
test_clear_follows_external_marker_when_include_dependent_dags_enabled(dag_maker,
session):
+ """With include_dependent_dags=True, clear() follows ExternalTaskMarker
links into child Dags."""
+ with dag_maker("parent_dag", session=session, schedule=None):
+ ExternalTaskMarker(
+ task_id="trigger_child",
+ external_dag_id="child_dag",
+ external_task_id="wait_for_parent",
+ recursion_depth=3,
+ )
+
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ serialized_parent = dag_maker.serialized_dag
+
+ with dag_maker("child_dag", session=session, schedule=None):
+ ExternalTaskSensor(
+ task_id="wait_for_parent",
+ external_dag_id="parent_dag",
+ external_task_id="trigger_child",
+ poke_interval=5,
+ )
+
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ session.flush()
+
+ result = serialized_parent.clear(
+ dry_run=True, only_failed=False, include_dependent_dags=True,
session=session
+ )
+
+ dag_ids = {ti.dag_id for ti in result}
+ task_ids = {ti.task_id for ti in result}
+
+ assert "parent_dag" in dag_ids
+ assert "child_dag" in dag_ids
+ assert "wait_for_parent" in task_ids
+
+
+def test_clear_reuses_provided_dag_bag_for_external_dags(dag_maker, session):
+ """Passing an existing dag_bag into clear() reuses it instead of creating
an uncached, un-configured one."""
+ with dag_maker("parent_dag", session=session, schedule=None):
+ ExternalTaskMarker(
+ task_id="trigger_child",
+ external_dag_id="child_dag",
+ external_task_id="wait_for_parent",
+ recursion_depth=3,
+ )
+
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ serialized_parent = dag_maker.serialized_dag
+
+ with dag_maker("child_dag", session=session, schedule=None):
+ ExternalTaskSensor(
+ task_id="wait_for_parent",
+ external_dag_id="parent_dag",
+ external_task_id="trigger_child",
+ poke_interval=5,
+ )
+
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ session.flush()
+
+ provided_dag_bag = DBDagBag()
+
+ with mock.patch("airflow.serialization.definitions.dag.DBDagBag",
wraps=DBDagBag) as mock_dbdagbag_cls:
+ result = serialized_parent.clear(
+ dry_run=True,
+ only_failed=False,
+ include_dependent_dags=True,
+ session=session,
+ dag_bag=provided_dag_bag,
+ )
+
+ mock_dbdagbag_cls.assert_not_called()
+ dag_ids = {ti.dag_id for ti in result}
+ assert "child_dag" in dag_ids
+
+
+def test_clear_creates_dag_bag_when_none_provided(dag_maker, session):
+ """Without a caller-provided dag_bag, clear() falls back to creating its
own."""
+ with dag_maker("parent_dag", session=session, schedule=None):
+ ExternalTaskMarker(
+ task_id="trigger_child",
+ external_dag_id="child_dag",
+ external_task_id="wait_for_parent",
+ recursion_depth=3,
+ )
+
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ serialized_parent = dag_maker.serialized_dag
+
+ with dag_maker("child_dag", session=session, schedule=None):
+ ExternalTaskSensor(
+ task_id="wait_for_parent",
+ external_dag_id="parent_dag",
+ external_task_id="trigger_child",
+ poke_interval=5,
+ )
+
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ session.flush()
+
+ with mock.patch("airflow.serialization.definitions.dag.DBDagBag",
wraps=DBDagBag) as mock_dbdagbag_cls:
+ serialized_parent.clear(dry_run=True, only_failed=False,
include_dependent_dags=True, session=session)
+
+ mock_dbdagbag_cls.assert_called_once_with(load_op_links=False)
+
+
+def
test_clear_dependent_dags_deserializes_child_dag_once_across_multiple_markers(dag_maker,
session):
+ """Multiple ExternalTaskMarkers into the same child dag must not each
re-deserialize it."""
+ with dag_maker("parent_dag", session=session, schedule=None):
+ ExternalTaskMarker(
+ task_id="trigger_child_a",
+ external_dag_id="child_dag",
+ external_task_id="wait_for_parent_a",
+ recursion_depth=3,
+ )
+ ExternalTaskMarker(
+ task_id="trigger_child_b",
+ external_dag_id="child_dag",
+ external_task_id="wait_for_parent_b",
+ recursion_depth=3,
+ )
+
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ serialized_parent = dag_maker.serialized_dag
+
+ with dag_maker("child_dag", session=session, schedule=None):
+ ExternalTaskSensor(
+ task_id="wait_for_parent_a",
+ external_dag_id="parent_dag",
+ external_task_id="trigger_child_a",
+ poke_interval=5,
+ )
+ ExternalTaskSensor(
+ task_id="wait_for_parent_b",
+ external_dag_id="parent_dag",
+ external_task_id="trigger_child_b",
+ poke_interval=5,
+ )
+
+ dag_maker.create_dagrun(logical_date=EXTERNAL_LOGICAL_DATE)
+ session.flush()
+
+ with mock.patch.object(
+ DBDagBag, "_read_dag", autospec=True, side_effect=DBDagBag._read_dag
+ ) as mock_read_dag:
+ result = serialized_parent.clear(
+ dry_run=True, only_failed=False, include_dependent_dags=True,
session=session
+ )
+
+ dag_ids = {ti.dag_id for ti in result}
+ assert "child_dag" in dag_ids
+ # Only the first ExternalTaskMarker into child_dag triggers an actual
deserialize; the
+ # second is served from the shared DBDagBag cache instead of
re-reading/re-deserializing.
+ child_dag_reads = [call for call in mock_read_dag.call_args_list if
call.args[1].dag_id == "child_dag"]
+ assert len(child_dag_reads) == 1
+
+
+def test_clear_raises_when_recursion_depth_exceeded(dag_maker, session):
+ """AirflowException is raised when the dependency chain depth exceeds
recursion_depth."""
Review Comment:
Should it be MaxRecursionDepthError instead?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
