jroachgolf84 commented on code in PR #65314:
URL: https://github.com/apache/airflow/pull/65314#discussion_r3606948056
##########
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:
Doesn't it feel quite heavy to thread the entire `DagBag` through?
--
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]
