kaxil commented on code in PR #65314:
URL: https://github.com/apache/airflow/pull/65314#discussion_r3605628549
##########
airflow-core/src/airflow/serialization/definitions/dag.py:
##########
@@ -1104,6 +1124,111 @@ def apply_state_filter(query):
else:
tis_full = apply_state_filter(tis_full)
+ if include_dependent_dags:
+ # Recursively find external tasks indicated by ExternalTaskMarker
+ 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",
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."
+ )
+
+ if ti.dag_run.logical_date is None:
+ continue
+
+ # Retrieve the logical date from the TI, Dag/task ID's
from the task
+ logical_date_str: str = ti.dag_run.logical_date.isoformat()
+ external_dag_id = task.external_dag_id
+ external_task_id = task.external_task_id
+
+ rendered =
RenderedTaskInstanceFields.get_templated_fields(ti, session=session)
+
+ if rendered:
+ external_dag_id = rendered.get("external_dag_id",
external_dag_id)
+ external_task_id = rendered.get("external_task_id",
external_task_id)
+
+ if "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 == external_dag_id,
+ TaskInstance.task_id == external_task_id,
+ DagRun.logical_date == external_logical_date,
+ )
+ )
+
+ # Load the DagBag such that Dags can be extracted from it
+ if not dag_bag:
+ dag_bag = DBDagBag(load_op_links=False)
Review Comment:
This is the first spot in the API server that builds a
`DBDagBag(load_op_links=False)`. Loading a child DAG runs
`serialized_dag.py:899`, which sets the process-global
`DagSerialization._load_operator_extra_links = self.load_op_links` and never
restores it. Since this endpoint is a sync `def`, it runs in the threadpool, so
a concurrent Graph/Grid/details request deserializing a DAG can read the
flipped flag and drop operator extra links until the next serial deserialize
resets it. The `_load_operator_extra_links = True` you added to
`test_dag_serialization.py`'s setup is that same leak showing up in tests.
Could we thread the request's existing `dag_bag` (already `load_op_links=True`)
into `_get_task_instances`, or save/restore the flag around this usage?
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -925,30 +929,48 @@ 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
+
task_instances: Sequence[TI]
- if dag_run_id is not None and not (past or future):
- # Use run_id-based clearing when we have a specific dag_run_id and not
using past/future
- task_instances = dag.clear(
- dry_run=True,
- task_ids=task_markers_to_clear,
- run_id=dag_run_id,
- session=session,
- run_on_latest_version=resolved_run_on_latest,
- only_failed=body.only_failed,
- only_running=body.only_running,
- )
- else:
- # Use date-based clearing when no dag_run_id or when past/future is
specified
- task_instances = dag.clear(
- dry_run=True,
- task_ids=task_markers_to_clear,
- start_date=body.start_date,
- end_date=body.end_date,
- session=session,
- run_on_latest_version=resolved_run_on_latest,
- only_failed=body.only_failed,
- only_running=body.only_running,
- )
+ try:
+ if dag_run_id is not None and not (past or future):
+ # Use run_id-based clearing when we have a specific dag_run_id and
not using past/future
+ task_instances = dag.clear(
+ dry_run=True,
+ task_ids=task_markers_to_clear,
+ run_id=dag_run_id,
+ session=session,
+ run_on_latest_version=resolved_run_on_latest,
+ only_failed=body.only_failed,
+ only_running=body.only_running,
+ include_dependent_dags=include_dependent_dags,
+ )
+ else:
+ # Use date-based clearing when no dag_run_id or when past/future
is specified
+ task_instances = dag.clear(
+ dry_run=True,
+ task_ids=task_markers_to_clear,
+ start_date=body.start_date,
+ end_date=body.end_date,
+ session=session,
+ run_on_latest_version=resolved_run_on_latest,
+ only_failed=body.only_failed,
+ only_running=body.only_running,
+ include_dependent_dags=include_dependent_dags,
+ )
+
+ except MaxRecursionDepthError 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
+
+ if include_dependent_dags:
+ # Ensure proper access to downstream dags/tasks with dag.clear and
include_dependent_dags
+ editable_dag_ids =
get_auth_manager().get_authorized_dag_ids(method="PUT", user=user)
Review Comment:
One asymmetry here: `get_authorized_dag_ids(method="PUT")` is a DAG-level
check, but the parent above gets authorized with `access_entity=TASK_INSTANCE`.
So on an auth manager that separates DAG-edit from task-instance-edit, the
child DAGs fall back to the coarser check. Not blocking -- it matches the
`get_authorized_dag_ids` precedent in import_error.py / dag_sources.py -- but
`batch_is_authorized_dag` with `access_entity=DagAccessEntity.TASK_INSTANCE`
per child would keep it in line with the parent.
##########
airflow-core/src/airflow/serialization/definitions/dag.py:
##########
@@ -1104,6 +1124,111 @@ def apply_state_filter(query):
else:
tis_full = apply_state_filter(tis_full)
+ if include_dependent_dags:
+ # Recursively find external tasks indicated by ExternalTaskMarker
+ 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",
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."
+ )
+
+ if ti.dag_run.logical_date is None:
+ continue
+
+ # Retrieve the logical date from the TI, Dag/task ID's
from the task
+ logical_date_str: str = ti.dag_run.logical_date.isoformat()
+ external_dag_id = task.external_dag_id
+ external_task_id = task.external_task_id
+
+ rendered =
RenderedTaskInstanceFields.get_templated_fields(ti, session=session)
+
+ if rendered:
+ external_dag_id = rendered.get("external_dag_id",
external_dag_id)
+ external_task_id = rendered.get("external_task_id",
external_task_id)
+
+ if "logical_date" in rendered:
+ logical_date_str = rendered["logical_date"]
+
+ external_logical_date = pendulum.parse(logical_date_str)
Review Comment:
Related to the RTIF edge case you're tracking for follow-up: if a custom
`logical_date` template renders to a non-ISO string, `pendulum.parse` raises
`ParserError`, which the route's `MaxRecursionDepthError`/`DagNotFound` handler
doesn't catch, so it comes back as a 500. Worth folding the parse failure into
that same follow-up as a 400.
--
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]
