bingqin2 commented on code in PR #73085:
URL: https://github.com/apache/airflow/pull/73085#discussion_r4006003276
##########
providers/standard/src/airflow/providers/standard/utils/sensor_helper.py:
##########
@@ -153,3 +156,78 @@ def _get_count_by_matched_states(
count += 1
return count
+
+
+def _check_external_task_existence(
+ api: Any,
+ *,
+ external_dag_id: str,
+ external_task_ids: Collection[str] | None,
+ external_task_group_id: str | None,
+ logical_dates: Collection[datetime] | None = None,
+ run_ids: Collection[str] | None = None,
+) -> bool:
+ """
+ Verify that the awaited tasks or task group exist in the awaited Dag runs,
through the execution API.
+
+ A Dag run's task instances are created together with the run, in the same
transaction and
+ from the run's own Dag version, so once a run exists its task instances
are the
+ version-accurate answer to whether a task belongs to it. Nothing can be
concluded about a run
+ that does not exist yet, which is why the caller has to repeat the check
until this function
+ returns True.
+
+ :param api: an object exposing ``get_dr_count``, ``get_ti_count`` and
``get_task_states`` the
+ way ``RuntimeTaskInstance`` does: the running task instance, or the
class itself.
+ :param external_dag_id: The ID of the external Dag.
+ :param external_task_ids: The task IDs that must exist in every awaited
run.
+ :param external_task_group_id: The task group ID that must exist in every
awaited run.
+ :param logical_dates: Logical dates identifying the awaited runs, used
when ``run_ids`` is empty.
+ :param run_ids: Run IDs identifying the awaited runs.
+ :return: True once every awaited run exists and passed the check, False
while at least one
+ awaited run does not exist yet.
+ :raises ExternalTaskNotFoundError: when an existing run has no task
instance for one of the tasks.
+ :raises ExternalTaskGroupNotFoundError: when the Dag has no such task
group, or an existing run
+ has no task instance for any task of the group.
+ """
+ from airflow.providers.standard.exceptions import (
+ ExternalTaskGroupNotFoundError,
+ ExternalTaskNotFoundError,
+ )
+ from airflow.sdk.exceptions import AirflowRuntimeError
+
+ awaited_runs: list[tuple[str, dict[str, list[Any]]]]
+ if run_ids:
+ awaited_runs = [(run_id, {"run_ids": [run_id]}) for run_id in run_ids]
+ else:
+ awaited_runs = [(dt.isoformat(), {"logical_dates": [dt]}) for dt in
logical_dates or []]
+
+ all_runs_checked = True
+ for run_label, run_filter in awaited_runs:
+ if api.get_dr_count(dag_id=external_dag_id, **run_filter) == 0:
+ all_runs_checked = False
+ continue
+ for task_id in external_task_ids or ():
+ if api.get_ti_count(dag_id=external_dag_id, task_ids=[task_id],
**run_filter) == 0:
Review Comment:
You're right, and thanks for the pointer to
`_verify_integrity_if_dag_changed`: for a run without a pinned bundle version
the scheduler adds newly parsed tasks on a later pass, so a zero count in that
window is not an answer, and the one-shot flag would have made it permanent.
That inference is gone. The check now asks `task-instances/states` for the
awaited task ids (or the task group) per existing run and treats only a 404 as
"does not exist"; a normal answer means the run's version defines it, whether
or not its instances exist yet. The 404 comes from #73086, which validates
`task_ids` against the version each named run resolves to and also counts an
existing instance as proof, so a task removed in a newer version is still found
for a run that has it. An API server without that validation never 404s for a
task, so with it the sensor simply keeps waiting, as it does today.
On mapped tasks: as far as I can see `DagRun._create_tasks` always creates a
`map_index=-1` placeholder while the length is unknown, so that case counted as
present, but with the count inference gone it no longer matters here.
##########
providers/standard/src/airflow/providers/standard/utils/sensor_helper.py:
##########
@@ -153,3 +156,78 @@ def _get_count_by_matched_states(
count += 1
return count
+
+
+def _check_external_task_existence(
+ api: Any,
+ *,
+ external_dag_id: str,
+ external_task_ids: Collection[str] | None,
+ external_task_group_id: str | None,
+ logical_dates: Collection[datetime] | None = None,
+ run_ids: Collection[str] | None = None,
+) -> bool:
+ """
+ Verify that the awaited tasks or task group exist in the awaited Dag runs,
through the execution API.
+
+ A Dag run's task instances are created together with the run, in the same
transaction and
+ from the run's own Dag version, so once a run exists its task instances
are the
+ version-accurate answer to whether a task belongs to it. Nothing can be
concluded about a run
+ that does not exist yet, which is why the caller has to repeat the check
until this function
+ returns True.
+
+ :param api: an object exposing ``get_dr_count``, ``get_ti_count`` and
``get_task_states`` the
+ way ``RuntimeTaskInstance`` does: the running task instance, or the
class itself.
+ :param external_dag_id: The ID of the external Dag.
+ :param external_task_ids: The task IDs that must exist in every awaited
run.
+ :param external_task_group_id: The task group ID that must exist in every
awaited run.
+ :param logical_dates: Logical dates identifying the awaited runs, used
when ``run_ids`` is empty.
+ :param run_ids: Run IDs identifying the awaited runs.
+ :return: True once every awaited run exists and passed the check, False
while at least one
+ awaited run does not exist yet.
+ :raises ExternalTaskNotFoundError: when an existing run has no task
instance for one of the tasks.
+ :raises ExternalTaskGroupNotFoundError: when the Dag has no such task
group, or an existing run
+ has no task instance for any task of the group.
+ """
+ from airflow.providers.standard.exceptions import (
+ ExternalTaskGroupNotFoundError,
+ ExternalTaskNotFoundError,
+ )
+ from airflow.sdk.exceptions import AirflowRuntimeError
+
+ awaited_runs: list[tuple[str, dict[str, list[Any]]]]
+ if run_ids:
+ awaited_runs = [(run_id, {"run_ids": [run_id]}) for run_id in run_ids]
+ else:
+ awaited_runs = [(dt.isoformat(), {"logical_dates": [dt]}) for dt in
logical_dates or []]
+
+ all_runs_checked = True
+ for run_label, run_filter in awaited_runs:
+ if api.get_dr_count(dag_id=external_dag_id, **run_filter) == 0:
+ all_runs_checked = False
+ continue
+ for task_id in external_task_ids or ():
+ if api.get_ti_count(dag_id=external_dag_id, task_ids=[task_id],
**run_filter) == 0:
+ raise ExternalTaskNotFoundError(
+ f"The external task {task_id} in Dag {external_dag_id}
does not exist for run {run_label}."
+ )
+
+ if external_task_group_id:
+ try:
+ run_id_task_state_map = api.get_task_states(
Review Comment:
Agreed. The 404 is now the only signal on the group path as well (an empty
result no longer raises), and the description states the dependency: with
#73086 the group is resolved against the awaited run's version, so a group
renamed after the run was created is still found for that run. Without it the
404 reflects the latest version, which is what `_poke_af3` already relies on
today: a renamed group makes the current sensor fail with an
`AirflowRuntimeError` on main, and this PR only changes the exception type. I
have converted this PR to a draft until #73086 lands.
--
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]
