kaxil commented on code in PR #75874:
URL: https://github.com/apache/airflow/pull/75874#discussion_r4516921186
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -593,6 +603,232 @@ 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
Review Comment:
I'd rather not. With targets stored in the Dag version, this would only help
runs on Dag versions serialized before the upgrade, and it brings back core
matching `ExternalTaskMarker` by name. The newsfragment says those runs don't
cascade. Ok to leave it?
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -593,6 +603,232 @@ 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)))
Review Comment:
Done, the exclusions are only computed for runs a target reaches.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -593,6 +603,232 @@ 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:
+ # One row per run, however many map indexes or archived attempts
the target task has.
+ target_runs = session.scalars(
+ select(DagRun).where(
+ DagRun.dag_id == target.dag_id,
+ DagRun.logical_date == target.logical_date,
+ select(TI.id)
+ .where(
+ TI.dag_id == DagRun.dag_id,
+ TI.run_id == DagRun.run_id,
+ TI.task_id == target.task_id,
+ TI.working_set.is_(True),
+ )
+ .exists(),
+ )
+ )
+ for target_run in target_runs:
+ if (dag := dag_bag.get_dag_for_run(target_run,
session=session)) is None:
+ log.warning(
+ "Not following an external clear target into a run
with no serialized Dag",
+ dag_id=target_run.dag_id,
+ run_id=target_run.run_id,
+ )
+ continue
+ task_ids = set(
+ dag.partial_subset(
+ task_ids=[target.task_id], include_downstream=True,
include_upstream=False
+ ).task_dict
+ )
+ run_key = (target_run.dag_id, target_run.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_run.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."""
Review Comment:
Applied.
##########
task-sdk/src/airflow/sdk/definitions/mappedoperator.py:
##########
@@ -255,7 +254,8 @@ def _expand(self, expand_input: ExpandInput, *, strict:
bool) -> MappedOperator:
template_fields_renderers=self.operator_class.template_fields_renderers,
ui_color=self.operator_class.ui_color,
ui_fgcolor=self.operator_class.ui_fgcolor,
- is_empty=issubclass(self.operator_class, EmptyOperator),
+ # Read the class attribute, like an unmapped operator does, so a
subclass can opt out.
Review Comment:
Applied.
--
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]
