uranusjr commented on code in PR #31642:
URL: https://github.com/apache/airflow/pull/31642#discussion_r1212771634
##########
airflow/utils/setup_teardown.py:
##########
@@ -110,25 +110,49 @@ def push_setup_teardown_task(cls, operator: Operator |
list[Operator]):
if first_task.is_teardown:
if not all(task.is_teardown == first_task.is_teardown for task
in operator):
raise ValueError("All tasks in the list must be either
setup or teardown tasks")
+ upstream_tasks = first_task.upstream_list
+ for task in upstream_tasks:
+ if not task.is_setup and not task.is_teardown:
+ raise ValueError(
+ "All upstream tasks in the context manager must be
a setup or teardown task"
+ )
SetupTeardownContext.push_context_managed_teardown_task(operator)
- upstream_setup: list[Operator] = [task for task in
first_task.upstream_list if task.is_setup]
+ upstream_setup: list[Operator] = [task for task in
upstream_tasks if task.is_setup]
if upstream_setup:
SetupTeardownContext.push_context_managed_setup_task(upstream_setup)
elif first_task.is_setup:
if not all(task.is_setup == first_task.is_setup for task in
operator):
raise ValueError("All tasks in the list must be either
setup or teardown tasks")
+ upstream_tasks = first_task.upstream_list
+ for task in upstream_tasks:
Review Comment:
```suggestion
for task in first_task.upstream_list:
```
Also for other blocks?
--
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]
