kaxil commented on code in PR #65751:
URL: https://github.com/apache/airflow/pull/65751#discussion_r3293446168
##########
task-sdk/src/airflow/sdk/bases/branch.py:
##########
@@ -61,13 +61,33 @@ def _expand_task_group_roots(
branches_to_execute = [branches_to_execute]
for branch in branches_to_execute:
- if branch in dag.task_group_dict:
- tg = dag.task_group_dict[branch]
+ resolved = self._resolve_branch_id(task, dag, branch)
+ if resolved in dag.task_group_dict:
+ tg = dag.task_group_dict[resolved]
root_ids = [root.task_id for root in tg.roots]
self.log.info("Expanding task group %s into %s", tg.group_id,
root_ids)
yield from root_ids
else:
- yield branch
+ yield resolved
+
+ def _resolve_branch_id(self, task, dag, branch: str) -> str:
+ """
+ Resolve a branch id, allowing relative ids when called from inside a
task group.
+
+ If ``branch`` matches a task or task group id as-is, return it
unchanged.
+ Otherwise, if the caller is inside a prefixed task group, try resolving
+ ``branch`` as ``"{group_id}.{branch}"`` so a branch inside group ``g``
+ can return ``"path_a"`` instead of ``"g.path_a"``.
+ """
Review Comment:
Worth documenting the lookup precedence: if a top-level `path_a` and a
`g.path_a` both exist, the bare `path_a` returned from inside `g` resolves to
the top-level task (the first `branch in dag.task_dict` check returns before
relative resolution runs). Users with name collisions across groups could be
surprised by this.
Suggest appending one line to the docstring, e.g.: *"Absolute matches take
precedence: relative resolution only runs when the bare name does not already
match a top-level task or group."*
(Same applies to the copy in
`providers/standard/src/airflow/providers/standard/operators/branch.py`.)
--
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]
