kaxil commented on code in PR #65314:
URL: https://github.com/apache/airflow/pull/65314#discussion_r3482222478
##########
airflow-core/src/airflow/serialization/definitions/dag.py:
##########
@@ -1137,6 +1155,124 @@ def apply_state_filter(query):
else:
tis_full = apply_state_filter(tis_full)
+ if include_dependent_dags:
+ # Recursively find external tasks indicated by ExternalTaskMarker
+ import pendulum
+
+ from airflow.providers.standard.sensors.external_task import
ExternalTaskMarker
+
+ # Build a full-object query for identifying ExternalTaskMarker TIs
in the current set
+ if as_pk_tuple:
+ all_ti_rows = session.execute(tis_pk).all()
+ condition = TaskInstance.filter_for_tis(
+ TaskInstanceKey(**cols._mapping) for cols in all_ti_rows
+ )
+ marker_query = select(TaskInstance).where(condition) if
condition is not None else None
+ else:
+ marker_query = tis_full
+
+ if marker_query is not None:
+ if visited_external_tis is None:
+ visited_external_tis = set()
+
+ external_marker_tis = session.scalars(
+ marker_query.where(TaskInstance.operator ==
ExternalTaskMarker.__name__)
+ )
+
+ for ti in external_marker_tis:
+ ti_key = ti.key.primary
+
+ if ti_key in visited_external_tis:
+ continue
+
+ visited_external_tis.add(ti_key)
+ task: ExternalTaskMarker = cast(
+ "ExternalTaskMarker",
copy.copy(self.get_task(ti.task_id))
+ )
+
+ if max_recursion_depth is None:
+ # Maximum recursion depth is set from the first
ExternalTaskMarker encountered
+ max_recursion_depth = task.recursion_depth
+
+ if recursion_depth + 1 > max_recursion_depth:
+ raise MaxRecursionDepthError(
Review Comment:
`MaxRecursionDepthError` here, and `DagNotFound` at line 1249, both escape
through the dry-run `dag.clear()` in the route, which is outside the handler's
only try/except (that one catches `AirflowClearRunningTaskException`). Nothing
at the app level handles `AirflowException`/`DagNotFound` either, so a cyclic
or too-deep marker chain -- or a missing child DAG -- comes back as a 500
rather than a 4xx. Probably want to catch these in the route and return
400/404, plus a test for the cyclic case.
##########
airflow-core/src/airflow/serialization/definitions/dag.py:
##########
@@ -1137,6 +1155,124 @@ def apply_state_filter(query):
else:
tis_full = apply_state_filter(tis_full)
+ if include_dependent_dags:
+ # Recursively find external tasks indicated by ExternalTaskMarker
+ import pendulum
+
+ from airflow.providers.standard.sensors.external_task import
ExternalTaskMarker
+
+ # Build a full-object query for identifying ExternalTaskMarker TIs
in the current set
+ if as_pk_tuple:
+ all_ti_rows = session.execute(tis_pk).all()
+ condition = TaskInstance.filter_for_tis(
+ TaskInstanceKey(**cols._mapping) for cols in all_ti_rows
+ )
+ marker_query = select(TaskInstance).where(condition) if
condition is not None else None
+ else:
+ marker_query = tis_full
+
+ if marker_query is not None:
+ if visited_external_tis is None:
+ visited_external_tis = set()
+
+ external_marker_tis = session.scalars(
+ marker_query.where(TaskInstance.operator ==
ExternalTaskMarker.__name__)
+ )
+
+ for ti in external_marker_tis:
+ ti_key = ti.key.primary
+
+ if ti_key in visited_external_tis:
+ continue
+
+ visited_external_tis.add(ti_key)
+ task: ExternalTaskMarker = cast(
+ "ExternalTaskMarker",
copy.copy(self.get_task(ti.task_id))
+ )
+
+ if max_recursion_depth is None:
+ # Maximum recursion depth is set from the first
ExternalTaskMarker encountered
+ max_recursion_depth = task.recursion_depth
+
+ if recursion_depth + 1 > max_recursion_depth:
+ raise MaxRecursionDepthError(
+ f"Maximum recursion depth {max_recursion_depth}
reached for "
+ f"{ExternalTaskMarker.__name__} {ti.task_id}. "
+ f"Attempted to clear too many tasks or there may
be a cyclic dependency."
+ )
+
+ # Resolve the logical_date that the ExternalTaskMarker
points to
+ dr_logical_date = session.scalar(
+ select(DagRun.logical_date).where(
+ DagRun.dag_id == ti.dag_id, DagRun.run_id ==
ti.run_id
+ )
+ )
+
+ if dr_logical_date is None:
+ continue
+
+ logical_date_str: str = dr_logical_date.isoformat()
+
+ # Check whether a non-default template was used by
comparing the serialized template
+ # field on the task against the default "{{
logical_date.isoformat() }}" pattern
+ default_template = "{{ logical_date.isoformat() }}"
+
+ # If it differs, look up the rendered value from
RenderedTaskInstanceFields
+ if task.logical_date != default_template:
+ from airflow.models.renderedtifields import
RenderedTaskInstanceFields
+
+ rendered =
RenderedTaskInstanceFields.get_templated_fields(ti, session=session)
Review Comment:
The custom-template path has a case that isn't handled: if the marker uses a
non-default `logical_date` template and the RTIF row has already been pruned
(`num_dag_runs_to_retain_rendered_fields`), `rendered` comes back empty and
`logical_date_str` quietly stays as the marker run's own date from line 1214 --
so the child lookup targets the wrong DagRun. The test seeds the RTIF row, so
this case isn't exercised. Either fail loudly or add a test with the row
missing.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -922,6 +922,10 @@ 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
Review Comment:
This endpoint only authorizes the path `dag_id` --
`requires_access_dag(method="PUT", access_entity=TASK_INSTANCE)` at line 834.
With `include_dependent_dags` set, the recursion in `_get_task_instances` pulls
TIs out of the linked child DAGs (`task.external_dag_id`) and
`clear_task_instances` resets them, but nothing re-checks access on those
children. So someone who can only edit the parent ends up able to clear/re-run
tasks in child DAGs they can't otherwise touch.
Could we run the cleared child TIs through the caller's editable set
(`get_authorized_dag_ids(method="PUT", user=...)`), or re-check per distinct
child `dag_id` before the reset?
##########
airflow-core/src/airflow/serialization/definitions/dag.py:
##########
@@ -1137,6 +1155,124 @@ def apply_state_filter(query):
else:
tis_full = apply_state_filter(tis_full)
+ if include_dependent_dags:
+ # Recursively find external tasks indicated by ExternalTaskMarker
+ import pendulum
Review Comment:
Minor stuff while you're in here. `import pendulum` can live at the module
top rather than inline. The `copy.copy(self.get_task(...))` on line 1190
doesn't look necessary, since the task is only read and never mutated. Lines
1205-1209 re-query `DagRun.logical_date` even though `ti.dag_run.logical_date`
already has it. And the `DBDagBag` is being built once per external TI inside
the loop (1241-1244) when it could be constructed once above.
##########
airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -226,6 +226,12 @@ class ClearTaskInstancesBody(StrictBaseModel):
)
prevent_running_task: bool = False
note: Annotated[str, StringConstraints(max_length=1000)] | None = None
+ include_downstream_dags: bool = Field(
Review Comment:
Since this brings back user-visible clear behavior and adds a public API
field, it's a bit beyond a routine bugfix -- might be worth a newsfragment
(`airflow-core/newsfragments/65314.improvement.rst`).
Also, the route already cascades on `downstream=True`
(`include_dependent_dags = body.include_downstream_dags or downstream`), so
this new field overlaps with the existing `include_downstream`. A note in the
description on why both exist, and whether the public name should line up with
the internal `include_dependent_dags`, would help.
##########
airflow-core/src/airflow/serialization/definitions/dag.py:
##########
@@ -1137,6 +1155,124 @@ def apply_state_filter(query):
else:
tis_full = apply_state_filter(tis_full)
+ if include_dependent_dags:
+ # Recursively find external tasks indicated by ExternalTaskMarker
+ import pendulum
+
+ from airflow.providers.standard.sensors.external_task import
ExternalTaskMarker
+
+ # Build a full-object query for identifying ExternalTaskMarker TIs
in the current set
+ if as_pk_tuple:
+ all_ti_rows = session.execute(tis_pk).all()
+ condition = TaskInstance.filter_for_tis(
+ TaskInstanceKey(**cols._mapping) for cols in all_ti_rows
+ )
+ marker_query = select(TaskInstance).where(condition) if
condition is not None else None
+ else:
+ marker_query = tis_full
+
+ if marker_query is not None:
+ if visited_external_tis is None:
+ visited_external_tis = set()
+
+ external_marker_tis = session.scalars(
+ marker_query.where(TaskInstance.operator ==
ExternalTaskMarker.__name__)
+ )
+
+ for ti in external_marker_tis:
+ ti_key = ti.key.primary
+
+ if ti_key in visited_external_tis:
+ continue
+
+ visited_external_tis.add(ti_key)
+ task: ExternalTaskMarker = cast(
+ "ExternalTaskMarker",
copy.copy(self.get_task(ti.task_id))
+ )
+
+ if max_recursion_depth is None:
+ # Maximum recursion depth is set from the first
ExternalTaskMarker encountered
+ max_recursion_depth = task.recursion_depth
+
+ if recursion_depth + 1 > max_recursion_depth:
+ raise MaxRecursionDepthError(
+ f"Maximum recursion depth {max_recursion_depth}
reached for "
+ f"{ExternalTaskMarker.__name__} {ti.task_id}. "
+ f"Attempted to clear too many tasks or there may
be a cyclic dependency."
+ )
+
+ # Resolve the logical_date that the ExternalTaskMarker
points to
+ dr_logical_date = session.scalar(
+ select(DagRun.logical_date).where(
+ DagRun.dag_id == ti.dag_id, DagRun.run_id ==
ti.run_id
+ )
+ )
+
+ if dr_logical_date is None:
+ continue
+
+ logical_date_str: str = dr_logical_date.isoformat()
+
+ # Check whether a non-default template was used by
comparing the serialized template
+ # field on the task against the default "{{
logical_date.isoformat() }}" pattern
+ default_template = "{{ logical_date.isoformat() }}"
+
+ # If it differs, look up the rendered value from
RenderedTaskInstanceFields
+ if task.logical_date != default_template:
+ from airflow.models.renderedtifields import
RenderedTaskInstanceFields
+
+ rendered =
RenderedTaskInstanceFields.get_templated_fields(ti, session=session)
+
+ if rendered and "logical_date" in rendered:
+ logical_date_str = rendered["logical_date"]
+
+ external_logical_date = pendulum.parse(logical_date_str)
+ external_tis = session.scalars(
+ select(TaskInstance)
+ .join(TaskInstance.dag_run)
+ .where(
+ TaskInstance.dag_id == task.external_dag_id,
Review Comment:
`ExternalTaskMarker.template_fields` covers all three of `external_dag_id`,
`external_task_id` and `logical_date`, but only `logical_date` gets resolved
from the rendered fields above -- the other two are read straight off the
deserialized task here. If someone templates `external_dag_id` (say `"{{
var.value.child_dag }}"`), this filter compares against the literal Jinja
string and the child just silently doesn't get cleared.
The rendered values are already in the dict you fetch for `logical_date`
(`RenderedTaskInstanceFields.get_templated_fields(ti)`, the same call used at
`external_task.py:95-96`), so pulling all three from there would cover it.
--
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]
