uranusjr commented on code in PR #75874:
URL: https://github.com/apache/airflow/pull/75874#discussion_r4508195329
##########
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
Review Comment:
The record belongs to a single attempt. Clearing the marker creates a new
attempt with a new ID and no XCom. So a clear, followed by another clear before
the marker runs again, doesn't cascade the second time. Could the lookup fall
back to the latest successful earlier attempt of the same `(dag_id, run_id,
task_id, map_index)` when the current one has no record?
##########
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:
A marker task instance with no record is skipped silently. Can we add a
narrow fallback for task instances without a record? The fallback would apply
only when the task type is the standard provider's `ExternalTaskMarker`, both
IDs contain no Jinja, and `logical_date` is the old default `{{
logical_date.isoformat() }}` or None.
##########
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."""
Review Comment:
Do we need a mechanism to keep this in sync with the same constant in Task
SDK?
##########
providers/standard/src/airflow/providers/standard/sensors/external_task.py:
##########
@@ -680,6 +688,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:
The target is often known at parse time, so maybe we could avoid launching a
worker if there’s no templating happening at all? That removes the worker cost
for most markers, and also removes the one-attempt gap above for them, since
the target lives in the Dag version rather than in an attempt.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -2485,6 +2487,17 @@ def _render_map_index(context: Context, ti:
RuntimeTaskInstance, log: Logger) ->
return rendered_map_index
+def _record_external_clear_target(context: Context, ti: RuntimeTaskInstance,
log: Logger) -> None:
+ """Record the task's resolved cross-Dag clear target, so the API server
can follow it on clear."""
+ if not isinstance(ti.task, SupportsExternalClearTarget):
+ return
+ if (target := ti.task.get_external_clear_target(context)) is None:
+ return
+ log.debug("Recording external clear target", target=target)
+ # Straight to the metadata DB, like operator extra links: the API server
reads it there.
+ _xcom_push_to_db(ti, key=EXTERNAL_CLEAR_TARGET_XCOM_KEY,
value=target.serialize())
Review Comment:
We should try to block user somehow pushing the XCom key. There's a
precedent in `LOOP_XCOM_PREFIX`.
##########
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:
```suggestion
"""Return a check for whether ``user`` may clear task instances in a
Dag."""
```
##########
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:
This can probably be built lazily instead.
##########
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:
```suggestion
```
--
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]
