Fix waiting setup tasks when they are not a direct upstream

airflow/models/abstractoperator.py

@@ -293,6 +293,17 @@ def get_upstreams_only_setups_and_teardowns(self) -> Iterable[Operator]:
                     if t.is_teardown and not t == self:
                         yield t
 
+    def get_upstreams_only_setups(self) -> Iterable[Operator]:
+        """
+        Only upstream setups.
+
+        This method is meant to be used when we are checking task dependencies where we need
+        to wait for all the upstream setups to complete before we can run the task.
+        """
+        for task in self.get_upstreams_only_setups_and_teardowns():
+            if task.is_setup:
+                yield task
+
     def _iter_all_mapped_downstreams(self) -> Iterator[MappedOperator | MappedTaskGroup]:
         """Return mapped nodes that are direct dependencies of the current task.
 

airflow/models/dag.py

@@ -128,6 +128,7 @@
     with_row_locks,
 )
 from airflow.utils.state import DagRunState, TaskInstanceState
+from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.types import NOTSET, ArgNotSet, DagRunType, EdgeInfoType
 
 if TYPE_CHECKING:
@@ -717,6 +718,11 @@ def validate_setup_teardown(self):
         :meta private:
         """
         for task in self.tasks:
+            if task.is_setup:
+                for down_task in task.downstream_list:
+                    if not down_task.is_teardown and down_task.trigger_rule != TriggerRule.ALL_SUCCESS:
+                        # this is required to ensure consistent clearing behavior when upstream
+                        raise ValueError("Setup tasks must be followed with trigger rule ALL_SUCCESS.")
             FailStopDagInvalidTriggerRule.check(dag=self, trigger_rule=task.trigger_rule)
 
     def __repr__(self):

airflow/ti_deps/deps/trigger_rule_dep.py

@@ -20,10 +20,12 @@
 import collections
 import collections.abc
 import functools
-from typing import TYPE_CHECKING, Iterator, NamedTuple
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Iterator
 
 from sqlalchemy import and_, func, or_, select
 
+from airflow.models import MappedOperator
 from airflow.models.taskinstance import PAST_DEPENDS_MET
 from airflow.ti_deps.dep_context import DepContext
 from airflow.ti_deps.deps.base_ti_dep import BaseTIDep, TIDepStatus
@@ -34,10 +36,13 @@
     from sqlalchemy.orm import Session
     from sqlalchemy.sql.expression import ColumnOperators
 
+    from airflow.models.dagrun import DagRun
+    from airflow.models.operator import Operator
     from airflow.models.taskinstance import TaskInstance
 
 
-class _UpstreamTIStates(NamedTuple):
+@dataclass
+class _UpstreamTIStates:
     """States of the upstream tis for a specific ti.
 
     This is used to determine whether the specific ti can run in this iteration.
@@ -51,9 +56,12 @@ class _UpstreamTIStates(NamedTuple):
     done: int
     success_setup: int
     skipped_setup: int
+    failed_setup: int = 0
 
     @classmethod
-    def calculate(cls, finished_upstreams: Iterator[TaskInstance]) -> _UpstreamTIStates:
+    def calculate(
+        cls, finished_upstreams: Iterator[TaskInstance], finished_setup_upstream_tis
+    ) -> _UpstreamTIStates:
         """Calculate states for a task instance.
 
         ``counter`` is inclusive of ``setup_counter`` -- e.g. if there are 2 skipped upstreams, one
@@ -67,8 +75,9 @@ def calculate(cls, finished_upstreams: Iterator[TaskInstance]) -> _UpstreamTISta
         for ti in finished_upstreams:
             curr_state = {ti.state: 1}
             counter.update(curr_state)
-            if ti.task.is_setup:
-                setup_counter.update(curr_state)
+        for ti in finished_setup_upstream_tis:
+            curr_state = {ti.state: 1}
+            setup_counter.update(curr_state)
         return _UpstreamTIStates(
             success=counter.get(TaskInstanceState.SUCCESS, 0),
             skipped=counter.get(TaskInstanceState.SKIPPED, 0),
@@ -78,6 +87,8 @@ def calculate(cls, finished_upstreams: Iterator[TaskInstance]) -> _UpstreamTISta
             done=sum(counter.values()),
             success_setup=setup_counter.get(TaskInstanceState.SUCCESS, 0),
             skipped_setup=setup_counter.get(TaskInstanceState.SKIPPED, 0),
+            failed_setup=setup_counter.get(TaskInstanceState.FAILED, 0)
+            + setup_counter.get(TaskInstanceState.UPSTREAM_FAILED, 0),
         )
 
 
@@ -121,6 +132,13 @@ def _evaluate_trigger_rule(
         from airflow.models.operator import needs_expansion
         from airflow.models.taskinstance import TaskInstance
 
+        if ti.task.is_teardown:
+            setup_upstream_tasks = [task for task in ti.task.upstream_list if task.is_setup]
+        else:
+            setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
+
+        setup_upstream_tasks_ids = {task.task_id for task in setup_upstream_tasks}
+
         task = ti.task
         upstream_tasks = {t.task_id: t for t in task.upstream_list}
         trigger_rule = task.trigger_rule
@@ -152,6 +170,26 @@ def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None:
                 session=session,
             )
 
+        def _get_tis_count_from_tasks_list(dag_run: DagRun, tasks_ids_list: list[str]) -> int:
+            """Get the number of task instances from a list of tasks."""
+            return session.execute(
+                select(func.count("*"))
+                .where(TaskInstance.dag_id == dag_run.dag_id, TaskInstance.run_id == dag_run.run_id)
+                .where(TaskInstance.task_id.in_(tasks_ids_list))
+            ).scalar_one()
+
+        def _get_setup_upstream_tis_count(
+            current_ti: TaskInstance, setup_upstream_tasks_list: list[Operator]
+        ):
+            mapped_setup_upstream_ids = {
+                task.task_id
+                for task in setup_upstream_tasks_list
+                if isinstance(task, MappedOperator) or task.task_group != current_ti.task.task_group
+            }
+            return _get_tis_count_from_tasks_list(current_ti.dag_run, list(mapped_setup_upstream_ids)) + len(
+                [task for task in setup_upstream_tasks_list if task.task_id not in mapped_setup_upstream_ids]
+            )
+
         def _is_relevant_upstream(upstream: TaskInstance) -> bool:
             """Whether a task instance is a "relevant upstream" of the current task."""
             # Not actually an upstream task.
@@ -176,12 +214,32 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool:
                 return True
             return False
 
+        def _is_relevant_setup_upstream(upstream: TaskInstance) -> bool:
+            """Whether a task instance is a "relevant upstream setup" of the current task."""
+            if upstream.task_id in setup_upstream_tasks_ids:
+                relevant = _get_relevant_upstream_map_indexes(upstream.task_id)
+                if relevant is None:
+                    return True
+                if relevant == upstream.map_index:
+                    return True
+                if isinstance(relevant, collections.abc.Container) and upstream.map_index in relevant:
+                    return True
+                return False
+            return False
+
         finished_upstream_tis = (
             finished_ti
             for finished_ti in dep_context.ensure_finished_tis(ti.get_dagrun(session), session)
             if _is_relevant_upstream(finished_ti)
         )
-        upstream_states = _UpstreamTIStates.calculate(finished_upstream_tis)
+
+        finished_setup_upstream_tis = (
+            finished_ti
+            for finished_ti in dep_context.ensure_finished_tis(ti.get_dagrun(session), session)
+            if _is_relevant_setup_upstream(finished_ti)
+        )
+
+        upstream_states = _UpstreamTIStates.calculate(finished_upstream_tis, finished_setup_upstream_tis)
 
         success = upstream_states.success
         skipped = upstream_states.skipped
@@ -191,6 +249,7 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool:
         done = upstream_states.done
         success_setup = upstream_states.success_setup
         skipped_setup = upstream_states.skipped_setup
+        failed_setup = upstream_states.failed_setup
 
         def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
             # Optimization: If the current task is not in a mapped task group,
@@ -223,9 +282,12 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
 
         # Optimization: Don't need to hit the database if all upstreams are
         # "simple" tasks (no task or task group mapping involved).
+        if not ti.task.is_teardown:
+            upstream_setup = _get_setup_upstream_tis_count(ti, setup_upstream_tasks)
         if not any(needs_expansion(t) for t in upstream_tasks.values()):
             upstream = len(upstream_tasks)
-            upstream_setup = sum(1 for x in upstream_tasks.values() if x.is_setup)
+            if ti.task.is_teardown:
+                upstream_setup = sum(1 for x in upstream_tasks.values() if x.is_setup)
         else:
             task_id_counts = session.execute(
                 select(TaskInstance.task_id, func.count(TaskInstance.task_id))
@@ -234,59 +296,72 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
                 .group_by(TaskInstance.task_id)
             ).all()
             upstream = sum(count for _, count in task_id_counts)
-            upstream_setup = sum(c for t, c in task_id_counts if upstream_tasks[t].is_setup)
+            if ti.task.is_teardown:
+                upstream_setup = sum(c for t, c in task_id_counts if upstream_tasks[t].is_setup)
 
         upstream_done = done >= upstream
+        setup_done = (success_setup + skipped_setup + failed_setup) >= upstream_setup
 
         changed = False
         new_state = None
         if dep_context.flag_upstream_failed:
-            if trigger_rule == TR.ALL_SUCCESS:
-                if upstream_failed or failed:
-                    new_state = TaskInstanceState.UPSTREAM_FAILED
-                elif skipped:
-                    new_state = TaskInstanceState.SKIPPED
-                elif removed and success and ti.map_index > -1:
-                    if ti.map_index >= success:
-                        new_state = TaskInstanceState.REMOVED
-            elif trigger_rule == TR.ALL_FAILED:
-                if success or skipped:
-                    new_state = TaskInstanceState.SKIPPED
-            elif trigger_rule == TR.ONE_SUCCESS:
-                if upstream_done and done == skipped:
-                    # if upstream is done and all are skipped mark as skipped
-                    new_state = TaskInstanceState.SKIPPED
-                elif upstream_done and success <= 0:
-                    # if upstream is done and there are no success mark as upstream failed
-                    new_state = TaskInstanceState.UPSTREAM_FAILED
-            elif trigger_rule == TR.ONE_FAILED:
-                if upstream_done and not (failed or upstream_failed):
-                    new_state = TaskInstanceState.SKIPPED
-            elif trigger_rule == TR.ONE_DONE:
-                if upstream_done and not (failed or success):
-                    new_state = TaskInstanceState.SKIPPED
-            elif trigger_rule == TR.NONE_FAILED:
-                if upstream_failed or failed:
-                    new_state = TaskInstanceState.UPSTREAM_FAILED
-            elif trigger_rule == TR.NONE_FAILED_MIN_ONE_SUCCESS:
-                if upstream_failed or failed:
-                    new_state = TaskInstanceState.UPSTREAM_FAILED
-                elif skipped == upstream:
-                    new_state = TaskInstanceState.SKIPPED
-            elif trigger_rule == TR.NONE_SKIPPED:
-                if skipped:
-                    new_state = TaskInstanceState.SKIPPED
-            elif trigger_rule == TR.ALL_SKIPPED:
-                if success or failed:
-                    new_state = TaskInstanceState.SKIPPED
-            elif trigger_rule == TR.ALL_DONE_SETUP_SUCCESS:
-                if upstream_done and upstream_setup and skipped_setup >= upstream_setup:
-                    # when there is an upstream setup and they have all skipped, then skip
-                    new_state = TaskInstanceState.SKIPPED
-                elif upstream_done and upstream_setup and success_setup == 0:
-                    # when there is an upstream setup, if none succeeded, mark upstream failed
-                    # if at least one setup ran, we'll let it run
-                    new_state = TaskInstanceState.UPSTREAM_FAILED
+            if not task.is_teardown and failed_setup:
+                # we should exclude the teardown tasks from this check,
+                # because they should be run even if there is only one success setup task
+                new_state = TaskInstanceState.UPSTREAM_FAILED
+            elif not task.is_teardown and upstream_setup and setup_done and skipped_setup > 0:
+                # when there are upstream setup tasks and at least one of them is skipped, then skip
+                new_state = TaskInstanceState.SKIPPED
+            elif not upstream_setup or setup_done:
+                # if there are no upstream setup tasks or all of them are done,
+                # and we haven't set a new state, then we can check the upstream tasks
+                if trigger_rule == TR.ALL_SUCCESS:
+                    if upstream_failed or failed:
+                        new_state = TaskInstanceState.UPSTREAM_FAILED
+                    elif skipped:
+                        new_state = TaskInstanceState.SKIPPED
+                    elif removed and success and ti.map_index > -1:
+                        if ti.map_index >= success:
+                            new_state = TaskInstanceState.REMOVED
+                elif trigger_rule == TR.ALL_FAILED:
+                    if success or skipped:
+                        new_state = TaskInstanceState.SKIPPED
+                elif trigger_rule == TR.ONE_SUCCESS:
+                    if upstream_done and done == skipped:
+                        # if upstream is done and all are skipped mark as skipped
+                        new_state = TaskInstanceState.SKIPPED
+                    elif upstream_done and success <= 0:
+                        # if upstream is done and there are no success mark as upstream failed
+                        new_state = TaskInstanceState.UPSTREAM_FAILED
+                elif trigger_rule == TR.ONE_FAILED:
+                    if upstream_done and not (failed or upstream_failed):
+                        new_state = TaskInstanceState.SKIPPED
+                elif trigger_rule == TR.ONE_DONE:
+                    if upstream_done and not (failed or success):
+                        new_state = TaskInstanceState.SKIPPED
+                elif trigger_rule == TR.NONE_FAILED:
+                    if upstream_failed or failed:
+                        new_state = TaskInstanceState.UPSTREAM_FAILED
+                elif trigger_rule == TR.NONE_FAILED_MIN_ONE_SUCCESS:
+                    if upstream_failed or failed:
+                        new_state = TaskInstanceState.UPSTREAM_FAILED
+                    elif skipped == upstream:
+                        new_state = TaskInstanceState.SKIPPED
+                elif trigger_rule == TR.NONE_SKIPPED:
+                    if skipped:
+                        new_state = TaskInstanceState.SKIPPED
+                elif trigger_rule == TR.ALL_SKIPPED:
+                    if success or failed:
+                        new_state = TaskInstanceState.SKIPPED
+                elif trigger_rule == TR.ALL_DONE_SETUP_SUCCESS:
+                    if upstream_done and upstream_setup and skipped_setup >= upstream_setup:
+                        # when there is an upstream setup and they have all skipped, then skip
+                        new_state = TaskInstanceState.SKIPPED
+                    elif upstream_done and upstream_setup and setup_done and success_setup == 0:
+                        # when there is an upstream setup, if none succeeded, mark upstream failed
+                        # if at least one setup ran, we'll let it run
+                        new_state = TaskInstanceState.UPSTREAM_FAILED
+
         if new_state is not None:
             if new_state == TaskInstanceState.SKIPPED and dep_context.wait_for_past_depends_before_skipping:
                 past_depends_met = ti.xcom_pull(
@@ -302,6 +377,15 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
         if changed:
             dep_context.have_changed_ti_states = True
 
+        # first, check if all the setup upstream tasks are done, if not, we can't run this task
+        # we should exclude the teardown tasks from this check, because they should be run even
+        # if there is only one success setup task
+        if not task.is_teardown and (success_setup + skipped_setup) < upstream_setup:
+            yield self._failing_status(
+                reason=f"Waiting {upstream_setup - (success_setup + skipped_setup)}"
+                " setup task(s) to complete."
+            )
+
         if trigger_rule == TR.ONE_SUCCESS:
             if success <= 0:
                 yield self._failing_status(