shahar1 commented on code in PR #75874:
URL: https://github.com/apache/airflow/pull/75874#discussion_r4500760853
##########
providers/standard/src/airflow/providers/standard/sensors/external_task.py:
##########
@@ -680,6 +687,9 @@ class ExternalTaskMarker(EmptyOperator):
ui_color = "#4db7db"
operator_extra_links = [ExternalDagLink()]
+ # Where Airflow can follow a recorded clear target, the marker runs on a
worker so that its templates
+ # render there and the resolved target gets recorded, instead of being
marked successful unrun.
+ inherits_from_empty_operator = ExternalClearTarget is None
Review Comment:
**major:** once the marker runs on a worker, its default template fails in
any run with no logical date. `get_template_context` only sets `logical_date`
when the run has one (`task_runner.py:410`), and Dags default to
`StrictUndefined`, so `{{ logical_date.isoformat() }}` raises `UndefinedError:
'logical_date' is undefined`. I confirmed that render locally on `main`.
That covers asset-triggered runs (the scheduler passes `logical_date=None`),
REST triggers with a null `logical_date`, and partitioned UI triggers. Today
the marker is marked successful without rendering, so these runs pass. With
this change the marker fails and blocks its downstream tasks, even though the
Dag author didn't write a bad template.
A run with no logical date has nothing to point at by date, so one option is
a default template that tolerates the undefined value, with
`get_external_clear_target` returning `None` and the runner recording nothing.
A regression test for the no-logical-date case would pin it.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -593,6 +602,149 @@ def _patch_task_instance_note(
ti.task_instance_note.user_id = user.get_id()
+EXTERNAL_CLEAR_TARGET_XCOM_KEY = "_external_clear_target"
+"""XCom key the Task SDK's task runner records
``ExternalClearTarget.serialize()`` under."""
+
+
+class ExternalClearDepthExceeded(AirflowException):
+ """A clear followed more levels of cross-Dag clear targets than the
starting task allows."""
+
+
+class _ClearTarget(NamedTuple):
+ dag_id: str
+ task_id: str
+ logical_date: datetime
+ recursion_depth: int
+
+ @classmethod
+ def from_record(cls, data: dict[str, Any]) -> _ClearTarget:
+ return cls(
+ dag_id=data["dag_id"],
+ task_id=data["task_id"],
+ logical_date=datetime.fromisoformat(data["logical_date"]),
+ recursion_depth=data["recursion_depth"],
+ )
+
+
+def _recorded_targets(ti_ids: Iterable[UUID], *, session: Session) ->
list[_ClearTarget]:
+ read = build_xcom_read_query(
+ producer_ids=select(TI.id).where(TI.id.in_(ti_ids)),
key=EXTERNAL_CLEAR_TARGET_XCOM_KEY
+ )
+ rows = session.execute(read.with_only_columns(xcom_entity(read).value))
+ return [_ClearTarget.from_record(XComModel.deserialize_value(row)) for row
in rows]
Review Comment:
**major:** `XComModel.deserialize_value` runs `XComDecoder`, which rebuilds
any `__classname__` envelope, here inside the API server. The PR's premise is
that the API server can't run user code, and the Execution API accepts any
non-empty key, so any task can write this one. Locally on `main`, a
`builtins.tuple` envelope came back as a tuple. A class outside the allowlist
raised `ImportError`, which neither `deserialize_value`'s `except (ValueError,
TypeError)` nor the route catches, so that clear would return a 500. The
default allowlist is `airflow.*`, so I'm not claiming an exploit.
`deserialize_value_primitives` exists for this case. Its docstring says:
"The values it consumes are primitives, so XCom is read without reconstructing
objects in that process." (https://github.com/apache/airflow/pull/72718). The
record is a plain dict, so it should work as a direct swap.
Related: `_ClearTarget.from_record` trusts the value's shape, so a non-dict
or a missing key raises an uncaught `TypeError`/`KeyError`, and a bad date an
uncaught `ValueError`, because that call isn't inside the route's `except
ValueError`. Every clear with downstream from that task instance would then
return a 500. Rather than adding defaults, I'd validate the record and either
skip it with a warning or map it to a 4xx.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -593,6 +602,149 @@ def _patch_task_instance_note(
ti.task_instance_note.user_id = user.get_id()
+EXTERNAL_CLEAR_TARGET_XCOM_KEY = "_external_clear_target"
+"""XCom key the Task SDK's task runner records
``ExternalClearTarget.serialize()`` under."""
+
+
+class ExternalClearDepthExceeded(AirflowException):
+ """A clear followed more levels of cross-Dag clear targets than the
starting task allows."""
+
+
+class _ClearTarget(NamedTuple):
+ dag_id: str
+ task_id: str
+ logical_date: datetime
+ recursion_depth: int
+
+ @classmethod
+ def from_record(cls, data: dict[str, Any]) -> _ClearTarget:
+ return cls(
+ dag_id=data["dag_id"],
+ task_id=data["task_id"],
+ logical_date=datetime.fromisoformat(data["logical_date"]),
+ recursion_depth=data["recursion_depth"],
+ )
+
+
+def _recorded_targets(ti_ids: Iterable[UUID], *, session: Session) ->
list[_ClearTarget]:
+ read = build_xcom_read_query(
+ producer_ids=select(TI.id).where(TI.id.in_(ti_ids)),
key=EXTERNAL_CLEAR_TARGET_XCOM_KEY
+ )
+ rows = session.execute(read.with_only_columns(xcom_entity(read).value))
+ return [_ClearTarget.from_record(XComModel.deserialize_value(row)) for row
in rows]
+
+
+def find_external_clear_task_instances(
+ task_instances: Iterable[TI],
+ *,
+ dag_bag: DBDagBag,
+ only_failed: bool = False,
+ only_running: bool = False,
+ session: Session,
+) -> list[TI]:
+ """
+ Return the task instances in other Dags that clearing ``task_instances``
should also clear.
+
+ A task instance that recorded an :class:`~airflow.sdk.ExternalClearTarget`
when it ran pulls in the
+ target task and its downstream tasks, in the target Dag's run for that
logical date. Those can record
+ targets of their own, which are followed in turn. The first level's
largest ``recursion_depth`` caps how
+ many levels are followed, which stops a cycle of targets.
+
+ :raises ExternalClearDepthExceeded: when the targets go deeper than that
cap.
+ :raises DagNotFound: when a target's Dag run has no serialized Dag to read
its downstream tasks from.
+ """
+ seen = {ti.id for ti in task_instances}
+ frontier = set(seen)
+ found: list[TI] = []
+ max_depth: int | None = None
+ depth = 0
+ while frontier and (targets := _recorded_targets(frontier,
session=session)):
+ if max_depth is None:
+ max_depth = max(target.recursion_depth for target in targets)
+ depth += 1
+ if depth > max_depth:
+ raise ExternalClearDepthExceeded(
+ f"Clearing followed more than {max_depth} levels of external
clear targets, "
+ f"the last into Dag {targets[0].dag_id!r}. There may be a
cycle between Dags."
+ )
+ frontier = set()
+ for target in targets:
+ target_tis = session.scalars(
+ select(TI)
+ .join(TI.dag_run)
+ .where(
+ TI.dag_id == target.dag_id,
+ TI.task_id == target.task_id,
+ DagRun.logical_date == target.logical_date,
+ )
+ )
+ for target_ti in target_tis:
+ if (dag := dag_bag.get_dag_for_run(target_ti.dag_run,
session=session)) is None:
+ raise DagNotFound(f"Could not find Dag
{target_ti.dag_id!r}")
+ downstream = dag.partial_subset(
+ task_ids=[target_ti.task_id], include_downstream=True,
include_upstream=False
+ )
+ for ti in dag.clear(
+ task_ids=list(downstream.task_dict),
+ run_id=target_ti.run_id,
+ only_failed=only_failed,
+ only_running=only_running,
+ dry_run=True,
+ session=session,
+ ):
+ if ti.id not in seen:
+ seen.add(ti.id)
+ found.append(ti)
+ frontier.add(ti.id)
+ return found
+
+
+def _editable_task_instances(tis: list[TI], *, user: BaseUser, session:
Session) -> list[TI]:
+ """Drop the task instances of Dags the user may not clear; the route only
checked the path Dag."""
+ dag_ids = list({ti.dag_id for ti in tis})
+ dag_id_to_team = DagModel.get_dag_id_to_team_name_mapping(dag_ids,
session=session)
+ auth_manager = get_auth_manager()
+ editable = {
+ dag_id
+ for dag_id in dag_ids
+ if auth_manager.is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.TASK_INSTANCE,
+ details=DagDetails(id=dag_id,
team_name=dag_id_to_team.get(dag_id)),
+ user=user,
+ )
+ }
+ return [ti for ti in tis if ti.dag_id in editable]
+
+
+def _clear_dependent_task_instances(
+ tis: list[TI],
+ *,
+ body: ClearTaskInstancesBody,
+ dag_run_state: DagRunState | Literal[False],
+ session: Session,
+) -> list[TI]:
+ """Clear task instances in other Dags, each under its own Dag's
``rerun_with_latest_version``."""
+ lock_dag_runs(session, {(ti.dag_id, ti.run_id) for ti in tis},
refresh=True)
+ by_dag: dict[str, list[TI]] = defaultdict(list)
+ for ti in session.scalars(
+ select(TI).where(TI.id.in_([ti.id for ti in
tis])).execution_options(populate_existing=True)
+ ):
+ by_dag[ti.dag_id].append(ti)
+ whole_task = not body.only_failed and not body.only_running
+ cleared: list[TI] = []
+ for dag_id, dag_tis in by_dag.items():
+ cleared += clear_task_instances(
Review Comment:
**major:** this plain `clear_task_instances` call has no loop handling. The
route uses `apply_loop_clear_scope` for loop-aware Dags and
`dag.clear(dry_run=False)` uses `clear_task_instances_for_runs`, which archives
the later passes of a cleared gate. Here, `dag.clear(dry_run=True)` in
`find_external_clear_task_instances` also lists those later passes ("A dry run
also lists the task instances of later loop passes that clearing a gate
archives"), and this call then retries them in place. If a target Dag has a
loop gate downstream of the target task, the later passes come back alongside
the re-run gate instead of being archived.
This is from reading the code; I didn't run it. Going through the same
helper as `dag.clear` looks right, though `clear_task_instances_for_runs`
returns `None`, so the route would need another way to get the cleared list. A
test with a looping target Dag would pin it.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -593,6 +602,149 @@ def _patch_task_instance_note(
ti.task_instance_note.user_id = user.get_id()
+EXTERNAL_CLEAR_TARGET_XCOM_KEY = "_external_clear_target"
+"""XCom key the Task SDK's task runner records
``ExternalClearTarget.serialize()`` under."""
+
+
+class ExternalClearDepthExceeded(AirflowException):
+ """A clear followed more levels of cross-Dag clear targets than the
starting task allows."""
+
+
+class _ClearTarget(NamedTuple):
+ dag_id: str
+ task_id: str
+ logical_date: datetime
+ recursion_depth: int
+
+ @classmethod
+ def from_record(cls, data: dict[str, Any]) -> _ClearTarget:
+ return cls(
+ dag_id=data["dag_id"],
+ task_id=data["task_id"],
+ logical_date=datetime.fromisoformat(data["logical_date"]),
+ recursion_depth=data["recursion_depth"],
+ )
+
+
+def _recorded_targets(ti_ids: Iterable[UUID], *, session: Session) ->
list[_ClearTarget]:
Review Comment:
**nit:** `AGENTS.md` says: "Name functions and methods with action verbs:
`get_`, `extract_`, `find_`, `compute_`, `build_`, etc. Avoid noun-only names
like `_serialize_keys` or `_base_names`". That would make this
`_get_recorded_targets`, and `_editable_task_instances`
`_filter_editable_task_instances`.
##########
task-sdk/src/airflow/sdk/definitions/external_clear.py:
##########
@@ -0,0 +1,86 @@
+# 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.
+"""
+Cross-Dag clear propagation.
+
+An operator that defines ``get_external_clear_target(context)`` returning an
:class:`ExternalClearTarget`
+asks Airflow to also clear that task, and everything downstream of it,
whenever this task is cleared with
+its downstream tasks. The task runner calls the method after a successful run,
once templates have rendered,
+and records the result so the API server can follow it without running user
code.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+ from airflow.sdk.definitions.context import Context
+
+__all__ = ["EXTERNAL_CLEAR_TARGET_XCOM_KEY", "ExternalClearTarget",
"SupportsExternalClearTarget"]
+
+EXTERNAL_CLEAR_TARGET_XCOM_KEY = "_external_clear_target"
+"""XCom key the task runner records a resolved :class:`ExternalClearTarget`
under."""
+
+
+@runtime_checkable
+class SupportsExternalClearTarget(Protocol):
+ """An operator that names a task in another Dag to clear along with
itself."""
+
+ def get_external_clear_target(self, context: Context) ->
ExternalClearTarget: ...
+
+
+@dataclass(frozen=True)
+class ExternalClearTarget:
+ """
+ A task in another Dag that is cleared together with the task that declares
it.
+
+ :param dag_id: The Dag holding the task to clear.
+ :param task_id: The task to clear. Its downstream tasks in that Dag are
cleared with it.
+ :param logical_date: The logical date of the Dag run to clear in.
+ :param recursion_depth: How many levels of these links a clear follows
when it starts here, which
+ keeps a cycle of links from being followed forever.
+ """
+
+ dag_id: str
+ task_id: str
+ logical_date: datetime
+ recursion_depth: int = 10
+
+ def __post_init__(self) -> None:
+ if self.logical_date.tzinfo is None:
+ raise ValueError("logical_date must be timezone-aware")
+ if self.recursion_depth <= 0:
+ raise ValueError("recursion_depth must be a positive integer")
+
+ def serialize(self) -> dict[str, Any]:
+ return {
+ "dag_id": self.dag_id,
+ "task_id": self.task_id,
+ "logical_date": self.logical_date.isoformat(),
+ "recursion_depth": self.recursion_depth,
+ }
+
+ @classmethod
+ def deserialize(cls, data: dict[str, Any]) -> ExternalClearTarget:
Review Comment:
**nit:** nothing in the PR calls `deserialize`, since core parses the record
with `_ClearTarget.from_record`. Could it be dropped?
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -1269,6 +1273,27 @@ def select_candidates() -> Sequence[TI]:
.execution_options(populate_existing=True)
).all()
+ include_dependent_dags = (
+ body.include_downstream if body.include_dependent_dags is None else
body.include_dependent_dags
+ )
+ dependent_tis: list[TI] = []
+ if include_dependent_dags and task_instances:
+ try:
+ dependent_tis = find_external_clear_task_instances(
+ task_instances,
+ dag_bag=dag_bag,
+ only_failed=body.only_failed,
+ only_running=body.only_running,
+ session=session,
+ )
+ except ExternalClearDepthExceeded as e:
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, str(e)) from e
+ except DagNotFound as e:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, str(e)) from e
+ dependent_tis = _editable_task_instances(dependent_tis, user=user,
session=session)
Review Comment:
**minor:** this filter runs after traversal, so Dags the caller can't edit
still affect the outcome. A depth overrun or missing serialized Dag in one of
them aborts the whole clear, including the tasks the caller can edit, and the
depth error names that Dag (`services/public/task_instances.py:666`). The UI
also has no way to send `include_dependent_dags=false`, so a UI user hit by the
400 or 404 can't clear those tasks with downstream at all. That ties into
`bbovenzi`'s question on this PR.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -593,6 +602,149 @@ def _patch_task_instance_note(
ti.task_instance_note.user_id = user.get_id()
+EXTERNAL_CLEAR_TARGET_XCOM_KEY = "_external_clear_target"
+"""XCom key the Task SDK's task runner records
``ExternalClearTarget.serialize()`` under."""
+
+
+class ExternalClearDepthExceeded(AirflowException):
+ """A clear followed more levels of cross-Dag clear targets than the
starting task allows."""
+
+
+class _ClearTarget(NamedTuple):
+ dag_id: str
+ task_id: str
+ logical_date: datetime
+ recursion_depth: int
+
+ @classmethod
+ def from_record(cls, data: dict[str, Any]) -> _ClearTarget:
+ return cls(
+ dag_id=data["dag_id"],
+ task_id=data["task_id"],
+ logical_date=datetime.fromisoformat(data["logical_date"]),
+ recursion_depth=data["recursion_depth"],
+ )
+
+
+def _recorded_targets(ti_ids: Iterable[UUID], *, session: Session) ->
list[_ClearTarget]:
+ read = build_xcom_read_query(
+ producer_ids=select(TI.id).where(TI.id.in_(ti_ids)),
key=EXTERNAL_CLEAR_TARGET_XCOM_KEY
+ )
+ rows = session.execute(read.with_only_columns(xcom_entity(read).value))
+ return [_ClearTarget.from_record(XComModel.deserialize_value(row)) for row
in rows]
+
+
+def find_external_clear_task_instances(
+ task_instances: Iterable[TI],
+ *,
+ dag_bag: DBDagBag,
+ only_failed: bool = False,
+ only_running: bool = False,
+ session: Session,
+) -> list[TI]:
+ """
+ Return the task instances in other Dags that clearing ``task_instances``
should also clear.
+
+ A task instance that recorded an :class:`~airflow.sdk.ExternalClearTarget`
when it ran pulls in the
+ target task and its downstream tasks, in the target Dag's run for that
logical date. Those can record
+ targets of their own, which are followed in turn. The first level's
largest ``recursion_depth`` caps how
+ many levels are followed, which stops a cycle of targets.
+
+ :raises ExternalClearDepthExceeded: when the targets go deeper than that
cap.
+ :raises DagNotFound: when a target's Dag run has no serialized Dag to read
its downstream tasks from.
+ """
+ seen = {ti.id for ti in task_instances}
+ frontier = set(seen)
+ found: list[TI] = []
+ max_depth: int | None = None
+ depth = 0
+ while frontier and (targets := _recorded_targets(frontier,
session=session)):
+ if max_depth is None:
+ max_depth = max(target.recursion_depth for target in targets)
+ depth += 1
+ if depth > max_depth:
+ raise ExternalClearDepthExceeded(
+ f"Clearing followed more than {max_depth} levels of external
clear targets, "
+ f"the last into Dag {targets[0].dag_id!r}. There may be a
cycle between Dags."
+ )
+ frontier = set()
+ for target in targets:
+ target_tis = session.scalars(
+ select(TI)
+ .join(TI.dag_run)
+ .where(
+ TI.dag_id == target.dag_id,
+ TI.task_id == target.task_id,
+ DagRun.logical_date == target.logical_date,
+ )
+ )
+ for target_ti in target_tis:
+ if (dag := dag_bag.get_dag_for_run(target_ti.dag_run,
session=session)) is None:
+ raise DagNotFound(f"Could not find Dag
{target_ti.dag_id!r}")
+ downstream = dag.partial_subset(
+ task_ids=[target_ti.task_id], include_downstream=True,
include_upstream=False
+ )
+ for ti in dag.clear(
+ task_ids=list(downstream.task_dict),
+ run_id=target_ti.run_id,
+ only_failed=only_failed,
+ only_running=only_running,
+ dry_run=True,
+ session=session,
+ ):
+ if ti.id not in seen:
+ seen.add(ti.id)
+ found.append(ti)
+ frontier.add(ti.id)
+ return found
+
+
+def _editable_task_instances(tis: list[TI], *, user: BaseUser, session:
Session) -> list[TI]:
+ """Drop the task instances of Dags the user may not clear; the route only
checked the path Dag."""
+ dag_ids = list({ti.dag_id for ti in tis})
+ dag_id_to_team = DagModel.get_dag_id_to_team_name_mapping(dag_ids,
session=session)
+ auth_manager = get_auth_manager()
+ editable = {
+ dag_id
+ for dag_id in dag_ids
+ if auth_manager.is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.TASK_INSTANCE,
+ details=DagDetails(id=dag_id,
team_name=dag_id_to_team.get(dag_id)),
+ user=user,
+ )
+ }
+ return [ti for ti in tis if ti.dag_id in editable]
+
+
+def _clear_dependent_task_instances(
+ tis: list[TI],
+ *,
+ body: ClearTaskInstancesBody,
+ dag_run_state: DagRunState | Literal[False],
+ session: Session,
+) -> list[TI]:
+ """Clear task instances in other Dags, each under its own Dag's
``rerun_with_latest_version``."""
+ lock_dag_runs(session, {(ti.dag_id, ti.run_id) for ti in tis},
refresh=True)
+ by_dag: dict[str, list[TI]] = defaultdict(list)
+ for ti in session.scalars(
+ select(TI).where(TI.id.in_([ti.id for ti in
tis])).execution_options(populate_existing=True)
Review Comment:
**minor:** `dependent_tis` was selected before these locks were taken, and
this re-read only refreshes by ID. The route re-runs `select_candidates()`
under the lock for source tasks. Here, with `only_failed` or `only_running`, a
task instance whose state changed in between still gets cleared. Taking the
source and dependent locks in two batches also departs from the ordering
`lock_dag_runs` documents ("paths over overlapping runs acquire them in the
same order and cannot deadlock each other"), though a deadlock would need
concurrent clears across a cycle of Dags.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -1269,6 +1273,27 @@ def select_candidates() -> Sequence[TI]:
.execution_options(populate_existing=True)
).all()
+ include_dependent_dags = (
Review Comment:
**minor:** the route wiring has no tests. `code-review.instructions.md`
says: "Flag any changed or added behaviour without a corresponding test". The
new tests call `find_external_clear_task_instances` directly, so nothing covers
`include_dependent_dags` following `include_downstream`, the `false` opt-out,
dry-run merging, or the 400/404 mapping. `_editable_task_instances` and
`_clear_dependent_task_instances` have no tests either, and the first is an
authorization boundary, so it's the one I'd cover first.
##########
providers/standard/src/airflow/providers/standard/sensors/external_task.py:
##########
@@ -712,6 +722,21 @@ def __init__(
raise ValueError("recursion_depth should be a positive integer")
self.recursion_depth = recursion_depth
+ def get_external_clear_target(self, context: Context) ->
ExternalClearTarget:
+ error = f"logical_date must render to a date and time, got
{self.logical_date!r}"
+ try:
+ logical_date = pendulum.parse(self.logical_date)
Review Comment:
**nit:** in a Dag with `render_template_as_native_obj=True`, a template like
`{{ logical_date }}` should render to a `DateTime` rather than a string.
`pendulum.parse` raises `TypeError` on a `DateTime` (checked locally), which
skips the friendly error below, and a value that's already a datetime could be
used as it is.
--
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]
