shahar1 commented on code in PR #75874:
URL: https://github.com/apache/airflow/pull/75874#discussion_r4502774753
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -593,6 +603,237 @@ 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 _RecordedClearTarget(BaseModel):
+ """A clear target as a worker recorded it. Validated, since any task can
write the XCom key."""
+
+ model_config = ConfigDict(frozen=True)
+
+ dag_id: str
+ task_id: str
+ logical_date: AwareDatetime
+ recursion_depth: PositiveInt
+
+
+@dataclass(frozen=True)
+class DependentClear:
+ """What clearing pulls in from one run of another Dag."""
+
+ dag: SerializedDAG
+ run_id: str
+ task_ids: frozenset[str]
+ exclude_task_ids: frozenset[tuple[str, int]]
+ task_instances: list[TI]
+
+
+def _get_recorded_targets(ti_ids: Iterable[UUID], *, session: Session) ->
list[_RecordedClearTarget]:
+ read = build_xcom_read_query(
+ producer_ids=select(TI.id).where(TI.id.in_(ti_ids)),
key=EXTERNAL_CLEAR_TARGET_XCOM_KEY
+ )
+ entity = xcom_entity(read)
+ targets = []
+ for row in session.execute(read.with_only_columns(entity.task_instance_id,
entity.value)):
+ # Plain JSON types only: the API server must not rebuild objects from
a value a task wrote.
+ try:
+
targets.append(_RecordedClearTarget.model_validate(XComModel.deserialize_value_primitives(row)))
+ except ValidationError:
+ log.warning(
+ "Ignoring an unreadable external clear target",
task_instance_id=str(row.task_instance_id)
+ )
+ return targets
+
+
+def find_dependent_clears(
+ task_instances: Iterable[TI],
+ *,
+ dag_bag: DBDagBag,
+ can_clear: Callable[[str], bool],
+ only_failed: bool = False,
+ only_running: bool = False,
+ session: Session,
+) -> list[DependentClear]:
+ """
+ Find what clearing ``task_instances`` should also clear in other Dags.
+
+ 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. Targets in Dags
``can_clear`` rejects are not
+ followed, and neither is a target whose run has no serialized Dag to read
its downstream tasks from.
+
+ :raises ExternalClearDepthExceeded: when the targets go deeper than that
cap.
+ """
+ source = list(task_instances)
+ seen = {ti.id for ti in source}
+ # Executions the source selection already clears are left to it, so none
is cleared twice.
+ source_executions: defaultdict[tuple[str, str], set[tuple[str, int]]] =
defaultdict(set)
+ for ti in source:
+ source_executions[(ti.dag_id, ti.run_id)].add((ti.task_id,
get_public_map_index(ti, session=session)))
+ runs: dict[tuple[str, str], tuple[SerializedDAG, set[str]]] = {}
+ frontier = set(seen)
+ max_depth: int | None = None
+ depth = 0
+ while frontier:
+ targets = [t for t in _get_recorded_targets(frontier, session=session)
if can_clear(t.dag_id)]
+ if not targets:
+ break
+ 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(
Review Comment:
**nit:** this query has no `TI.working_set.is_(True)` filter, so every
archived attempt of the target task comes back as its own `target_ti`. The
`runs.setdefault` keeps the result right, but each extra row repeats
`get_dag_for_run` and the `dag.clear(dry_run=True)` pass for the same run.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -593,6 +603,237 @@ 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 _RecordedClearTarget(BaseModel):
+ """A clear target as a worker recorded it. Validated, since any task can
write the XCom key."""
+
+ model_config = ConfigDict(frozen=True)
+
+ dag_id: str
+ task_id: str
+ logical_date: AwareDatetime
+ recursion_depth: PositiveInt
+
+
+@dataclass(frozen=True)
+class DependentClear:
+ """What clearing pulls in from one run of another Dag."""
+
+ dag: SerializedDAG
+ run_id: str
+ task_ids: frozenset[str]
+ exclude_task_ids: frozenset[tuple[str, int]]
+ task_instances: list[TI]
+
+
+def _get_recorded_targets(ti_ids: Iterable[UUID], *, session: Session) ->
list[_RecordedClearTarget]:
+ read = build_xcom_read_query(
+ producer_ids=select(TI.id).where(TI.id.in_(ti_ids)),
key=EXTERNAL_CLEAR_TARGET_XCOM_KEY
+ )
+ entity = xcom_entity(read)
+ targets = []
+ for row in session.execute(read.with_only_columns(entity.task_instance_id,
entity.value)):
+ # Plain JSON types only: the API server must not rebuild objects from
a value a task wrote.
+ try:
+
targets.append(_RecordedClearTarget.model_validate(XComModel.deserialize_value_primitives(row)))
+ except ValidationError:
+ log.warning(
+ "Ignoring an unreadable external clear target",
task_instance_id=str(row.task_instance_id)
+ )
+ return targets
+
+
+def find_dependent_clears(
+ task_instances: Iterable[TI],
+ *,
+ dag_bag: DBDagBag,
+ can_clear: Callable[[str], bool],
+ only_failed: bool = False,
+ only_running: bool = False,
+ session: Session,
+) -> list[DependentClear]:
+ """
+ Find what clearing ``task_instances`` should also clear in other Dags.
+
+ 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. Targets in Dags
``can_clear`` rejects are not
+ followed, and neither is a target whose run has no serialized Dag to read
its downstream tasks from.
+
+ :raises ExternalClearDepthExceeded: when the targets go deeper than that
cap.
+ """
+ source = list(task_instances)
+ seen = {ti.id for ti in source}
+ # Executions the source selection already clears are left to it, so none
is cleared twice.
+ source_executions: defaultdict[tuple[str, str], set[tuple[str, int]]] =
defaultdict(set)
+ for ti in source:
+ source_executions[(ti.dag_id, ti.run_id)].add((ti.task_id,
get_public_map_index(ti, session=session)))
+ runs: dict[tuple[str, str], tuple[SerializedDAG, set[str]]] = {}
+ frontier = set(seen)
+ max_depth: int | None = None
+ depth = 0
+ while frontier:
+ targets = [t for t in _get_recorded_targets(frontier, session=session)
if can_clear(t.dag_id)]
+ if not targets:
+ break
+ 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:
+ log.warning(
+ "Not following an external clear target into a run
with no serialized Dag",
+ dag_id=target_ti.dag_id,
+ run_id=target_ti.run_id,
+ )
+ continue
+ task_ids = set(
+ dag.partial_subset(
+ task_ids=[target_ti.task_id], include_downstream=True,
include_upstream=False
+ ).task_dict
+ )
+ run_key = (target_ti.dag_id, target_ti.run_id)
+ runs.setdefault(run_key, (dag, set()))[1].update(task_ids)
+ for ti in dag.clear(
+ task_ids=sorted(task_ids),
+ run_id=target_ti.run_id,
+ only_failed=only_failed,
+ only_running=only_running,
+ exclude_task_ids=frozenset(source_executions[run_key]),
+ dry_run=True,
+ session=session,
+ ):
+ if ti.id not in seen:
+ seen.add(ti.id)
+ frontier.add(ti.id)
+ clears = []
+ for (dag_id, run_id), (dag, task_ids) in runs.items():
+ exclude_task_ids = frozenset(source_executions[(dag_id, run_id)])
+ tis = list(
+ dag.clear(
+ task_ids=sorted(task_ids),
+ run_id=run_id,
+ only_failed=only_failed,
+ only_running=only_running,
+ exclude_task_ids=exclude_task_ids,
+ dry_run=True,
+ session=session,
+ )
+ )
+ if tis:
+ clears.append(
+ DependentClear(
+ dag=dag,
+ run_id=run_id,
+ task_ids=frozenset(task_ids),
+ exclude_task_ids=exclude_task_ids,
+ task_instances=tis,
+ )
+ )
+ return clears
+
+
+def build_dag_clear_check(user: BaseUser, *, session: Session) ->
Callable[[str], bool]:
+ """Return a check for whether ``user`` may clear task instances in a Dag;
the route only checked one."""
+ auth_manager = get_auth_manager()
+ decisions: dict[str, bool] = {}
+
+ def can_clear(dag_id: str) -> bool:
+ if dag_id not in decisions:
+ decisions[dag_id] = auth_manager.is_authorized_dag(
+ method="PUT",
+ access_entity=DagAccessEntity.TASK_INSTANCE,
+ details=DagDetails(id=dag_id,
team_name=DagModel.get_team_name(dag_id, session=session)),
+ user=user,
+ )
+ return decisions[dag_id]
+
+ return can_clear
+
+
+def clear_dependent_runs(
+ clears: list[DependentClear],
+ *,
+ body: ClearTaskInstancesBody,
+ dag_run_state: DagRunState | Literal[False],
+ session: Session,
+) -> list[TI]:
+ """
+ Clear what :func:`find_dependent_clears` found, selecting it again under
the run locks.
+
+ Each Dag clears under its own ``rerun_with_latest_version``, and
``SerializedDAG.clear`` archives the
+ later passes of any loop gate it clears.
+ """
+ lock_dag_runs(session, {(clear.dag.dag_id, clear.run_id) for clear in
clears}, refresh=True)
+ cleared: list[TI] = []
+ for clear in clears:
+ task_ids = sorted(clear.task_ids)
+ selected = {
Review Comment:
**nit:** this dry run repeats the one `find_dependent_clears` already did
under the same lock in the route, since both run after `lock_dag_runs`.
`clear.task_instances` holds that plan, so `selected` could be built from it
and save one traversal per dependent run. Fine to leave if you prefer the local
re-check.
--
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]
