From b7833dbbeb2f32da723ec17b4498de416156f43a Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Sun, 20 Aug 2023 22:23:22 +0200
Subject: [PATCH 01/22] Fix waiting setup tasks when they are not a direct
 upstream

---
 airflow/models/abstractoperator.py          |  10 +
 airflow/ti_deps/deps/trigger_rule_dep.py    |  48 ++-
 tests/models/test_taskinstance.py           | 366 +++++++++++++++-----
 tests/ti_deps/deps/test_trigger_rule_dep.py |  59 +++-
 4 files changed, 384 insertions(+), 99 deletions(-)

diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
index ba357c0bd1bbf..d3833e8f9b119 100644
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -293,6 +293,16 @@ 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.
+        """
+        upstream_setups = {x for x in self.get_flat_relatives(upstream=True) if x.is_setup}
+        yield from upstream_setups
+
     def _iter_all_mapped_downstreams(self) -> Iterator[MappedOperator | MappedTaskGroup]:
         """Return mapped nodes that are direct dependencies of the current task.
 
diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index c7e2982fffdd7..4a92dc2059c5d 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -34,6 +34,7 @@
     from sqlalchemy.orm import Session
     from sqlalchemy.sql.expression import ColumnOperators
 
+    from airflow.models import Operator
     from airflow.models.taskinstance import TaskInstance
 
 
@@ -51,6 +52,7 @@ class _UpstreamTIStates(NamedTuple):
     done: int
     success_setup: int
     skipped_setup: int
+    failed_setup: int
 
     @classmethod
     def calculate(cls, finished_upstreams: Iterator[TaskInstance]) -> _UpstreamTIStates:
@@ -78,6 +80,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),
         )
 
 
@@ -94,26 +98,33 @@ def _get_dep_statuses(
         session: Session,
         dep_context: DepContext,
     ) -> Iterator[TIDepStatus]:
+        # get all the setup tasks upstream of this task
+        setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
         # Checking that all upstream dependencies have succeeded.
-        if not ti.task.upstream_task_ids:
+        if not ti.task.upstream_task_ids and not setup_upstream_tasks:
             yield self._passing_status(reason="The task instance did not have any upstream tasks.")
             return
-        if ti.task.trigger_rule == TR.ALWAYS:
+        if ti.task.trigger_rule == TR.ALWAYS and not setup_upstream_tasks:
+            # even with ALWAYS trigger rule, we still need to check setup tasks
             yield self._passing_status(reason="The task had a always trigger rule set.")
             return
-        yield from self._evaluate_trigger_rule(ti=ti, dep_context=dep_context, session=session)
+        yield from self._evaluate_trigger_rule(
+            ti=ti, dep_context=dep_context, setup_upstream_tasks=setup_upstream_tasks, session=session
+        )
 
     def _evaluate_trigger_rule(
         self,
         *,
         ti: TaskInstance,
         dep_context: DepContext,
+        setup_upstream_tasks: list[Operator],
         session: Session,
     ) -> Iterator[TIDepStatus]:
         """Evaluate whether ``ti``'s trigger rule was met.
 
         :param ti: Task instance to evaluate the trigger rule of.
         :param dep_context: The current dependency context.
+        setup_upstream_tasks: The setup tasks upstream of the current task.
         :param session: Database session.
         """
         from airflow.models.abstractoperator import NotMapped
@@ -154,6 +165,9 @@ def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None:
 
         def _is_relevant_upstream(upstream: TaskInstance) -> bool:
             """Whether a task instance is a "relevant upstream" of the current task."""
+            # All the setup tasks upstreams are relevant event if they are not a direct upstream.
+            if upstream.task_id in map(lambda t: t.task_id, setup_upstream_tasks):
+                return True
             # Not actually an upstream task.
             if upstream.task_id not in task.upstream_task_ids:
                 return False
@@ -183,6 +197,8 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool:
         )
         upstream_states = _UpstreamTIStates.calculate(finished_upstream_tis)
 
+        upstream_setup = len(setup_upstream_tasks)  # count of setup tasks upstream of this task
+
         success = upstream_states.success
         skipped = upstream_states.skipped
         failed = upstream_states.failed
@@ -191,6 +207,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,
@@ -225,7 +242,6 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
         # "simple" tasks (no task or task group mapping involved).
         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)
         else:
             task_id_counts = session.execute(
                 select(TaskInstance.task_id, func.count(TaskInstance.task_id))
@@ -234,14 +250,19 @@ 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)
 
         upstream_done = done >= upstream
+        setup_done = (success_setup + skipped_setup + failed_setup) >= upstream_setup
+        is_tear_down = task.is_teardown or trigger_rule == TR.ALL_DONE_SETUP_SUCCESS
 
         changed = False
         new_state = None
         if dep_context.flag_upstream_failed:
-            if trigger_rule == TR.ALL_SUCCESS:
+            if not is_tear_down 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 trigger_rule == TR.ALL_SUCCESS:
                 if upstream_failed or failed:
                     new_state = TaskInstanceState.UPSTREAM_FAILED
                 elif skipped:
@@ -283,10 +304,14 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
                 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:
+                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 not is_tear_down and upstream_setup and not new_state 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
+
         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 +327,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 is_tear_down 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(
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index ac2f7f71748ef..eb0ccf7070fdd 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1103,59 +1103,116 @@ def test_depends_on_past(self, dag_maker):
     # Numeric fields are in order:
     #   successes, skipped, failed, upstream_failed, removed, done
     @pytest.mark.parametrize(
-        "trigger_rule, upstream_setups, upstream_states, flag_upstream_failed, expect_state, expect_passed",
+        "trigger_rule, direct_upstream_setups, indirect_upstream_setups, upstream_states,"
+        " flag_upstream_failed, expect_state, expect_passed",
         [
             #
             # Tests for all_success
             #
-            ["all_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
-            ["all_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
-            ["all_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["all_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
+            ["all_success", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_success", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
+            [
+                "all_success",
+                0,
+                0,
+                _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0),
+                True,
+                State.UPSTREAM_FAILED,
+                False,
+            ],
+            ["all_success", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
             #
             # Tests for one_success
             #
-            ["one_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
-            ["one_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, True],
-            ["one_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, True],
-            ["one_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, True],
-            ["one_success", 0, _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
-            ["one_success", 0, _UpstreamTIStates(0, 4, 1, 0, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", 0, _UpstreamTIStates(0, 3, 1, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", 0, _UpstreamTIStates(0, 4, 0, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", 0, _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", 0, _UpstreamTIStates(0, 0, 4, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", 0, _UpstreamTIStates(0, 0, 0, 5, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["one_success", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, True],
+            ["one_success", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
+            ["one_success", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, True],
+            ["one_success", 0, 0, _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            [
+                "one_success",
+                0,
+                0,
+                _UpstreamTIStates(0, 4, 1, 0, 0, 5, 0, 0, 0),
+                True,
+                State.UPSTREAM_FAILED,
+                False,
+            ],
+            [
+                "one_success",
+                0,
+                0,
+                _UpstreamTIStates(0, 3, 1, 1, 0, 5, 0, 0, 0),
+                True,
+                State.UPSTREAM_FAILED,
+                False,
+            ],
+            [
+                "one_success",
+                0,
+                0,
+                _UpstreamTIStates(0, 4, 0, 1, 0, 5, 0, 0, 0),
+                True,
+                State.UPSTREAM_FAILED,
+                False,
+            ],
+            [
+                "one_success",
+                0,
+                0,
+                _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0),
+                True,
+                State.UPSTREAM_FAILED,
+                False,
+            ],
+            [
+                "one_success",
+                0,
+                0,
+                _UpstreamTIStates(0, 0, 4, 1, 0, 5, 0, 0, 0),
+                True,
+                State.UPSTREAM_FAILED,
+                False,
+            ],
+            [
+                "one_success",
+                0,
+                0,
+                _UpstreamTIStates(0, 0, 0, 5, 0, 5, 0, 0, 0),
+                True,
+                State.UPSTREAM_FAILED,
+                False,
+            ],
             #
             # Tests for all_failed
             #
-            ["all_failed", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", 0, _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0), True, None, True],
-            ["all_failed", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, 0, _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_failed", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
             #
             # Tests for one_failed
             #
-            ["one_failed", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
-            ["one_failed", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
-            ["one_failed", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, True],
-            ["one_failed", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
-            ["one_failed", 0, _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
+            ["one_failed", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["one_failed", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
+            ["one_failed", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
+            ["one_failed", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["one_failed", 0, 0, _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
             #
             # Tests for done
             #
-            ["all_done", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
-            ["all_done", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
-            ["all_done", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, False],
-            ["all_done", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
+            ["all_done", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_done", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
+            ["all_done", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
             #
             # Tests for all_done_setup_success: no upstream setups -> same as all_done
             #
-            ["all_done_setup_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
-            ["all_done_setup_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
-            ["all_done_setup_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, False],
-            ["all_done_setup_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
+            ["all_done_setup_success", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_done_setup_success", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
+            ["all_done_setup_success", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done_setup_success", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
             #
             # Tests for all_done_setup_success: with upstream setups -> different from all_done
             #
@@ -1171,7 +1228,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(6, 0, 0, 0, 0, 6, 1, 0),
+                0,
+                _UpstreamTIStates(6, 0, 0, 0, 0, 6, 1, 0, 0),
                 True,
                 None,
                 True,
@@ -1180,7 +1238,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 2, 0),
+                0,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 2, 0, 0),
                 True,
                 None,
                 True,
@@ -1189,7 +1248,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(5, 0, 1, 0, 0, 6, 0, 0),
+                0,
+                _UpstreamTIStates(5, 0, 1, 0, 0, 6, 0, 0, 1),
                 True,
                 State.UPSTREAM_FAILED,
                 False,
@@ -1198,7 +1258,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                _UpstreamTIStates(5, 0, 2, 0, 0, 7, 0, 0),
+                0,
+                _UpstreamTIStates(5, 0, 2, 0, 0, 7, 0, 0, 2),
                 True,
                 State.UPSTREAM_FAILED,
                 False,
@@ -1207,7 +1268,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(5, 1, 0, 0, 0, 6, 0, 1),
+                0,
+                _UpstreamTIStates(5, 1, 0, 0, 0, 6, 0, 1, 0),
                 True,
                 State.SKIPPED,
                 False,
@@ -1216,7 +1278,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                _UpstreamTIStates(5, 2, 0, 0, 0, 7, 0, 2),
+                0,
+                _UpstreamTIStates(5, 2, 0, 0, 0, 7, 0, 2, 0),
                 True,
                 State.SKIPPED,
                 False,
@@ -1225,7 +1288,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                _UpstreamTIStates(5, 1, 1, 0, 0, 7, 0, 1),
+                0,
+                _UpstreamTIStates(5, 1, 1, 0, 0, 7, 0, 1, 1),
                 True,
                 State.UPSTREAM_FAILED,
                 False,
@@ -1234,7 +1298,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                _UpstreamTIStates(6, 0, 1, 0, 0, 7, 1, 0),
+                0,
+                _UpstreamTIStates(6, 0, 1, 0, 0, 7, 1, 0, 1),
                 True,
                 None,
                 True,
@@ -1243,7 +1308,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                _UpstreamTIStates(6, 1, 0, 0, 0, 7, 1, 1),
+                0,
+                _UpstreamTIStates(6, 1, 0, 0, 0, 7, 1, 1, 0),
                 True,
                 None,
                 True,
@@ -1252,7 +1318,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(3, 0, 0, 0, 0, 3, 1, 0),
+                0,
+                _UpstreamTIStates(3, 0, 0, 0, 0, 3, 1, 0, 0),
                 True,
                 None,
                 False,
@@ -1261,7 +1328,8 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(3, 0, 1, 0, 0, 4, 1, 0),
+                0,
+                _UpstreamTIStates(3, 0, 1, 0, 0, 4, 1, 0, 0),
                 True,
                 None,
                 False,
@@ -1270,12 +1338,133 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(3, 1, 0, 0, 0, 4, 1, 0),
+                0,
+                _UpstreamTIStates(3, 1, 0, 0, 0, 4, 1, 0, 0),
                 True,
                 None,
                 False,
                 id="not all done, one skipped",
             ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 2, 0, 0),
+                True,
+                None,
+                True,
+                id="indirect upstream setups - all success",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 0, 0),
+                True,
+                None,
+                False,
+                id="indirect upstream setups - one not done and one success",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 1, 0),
+                True,
+                None,
+                False,
+                id="indirect upstream setups - one not done and one skipped",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 1, 0),
+                True,
+                TaskInstanceState.SKIPPED,
+                True,
+                id="indirect upstream setups - all done: one skipped and one success",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 0, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - all done: one success and one failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 1, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - all done: one skipped and one failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 0, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - one not done and one failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(0, 0, 1, 0, 0, 1, 2, 0, 0),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - all setup success but upstream failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(5, 0, 0, 0, 0, 5, 2, 0, 0),
+                True,
+                None,
+                False,
+                id="indirect upstream setups - all setup success but upstream still running",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - one setup failed but upstream still running",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(4, 0, 1, 0, 0, 5, 0, 2, 0),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - all setup skipped but upstream failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(0, 7, 0, 0, 0, 7, 0, 0, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - one setup failed but upstream skipped",
+            ),
         ],
     )
     def test_check_task_dependencies(
@@ -1283,7 +1472,8 @@ def test_check_task_dependencies(
         monkeypatch,
         dag_maker,
         trigger_rule: str,
-        upstream_setups: int,
+        direct_upstream_setups: int,
+        indirect_upstream_setups: int,
         upstream_states: _UpstreamTIStates,
         flag_upstream_failed: bool,
         expect_state: State,
@@ -1293,18 +1483,26 @@ def test_check_task_dependencies(
 
         # sanity checks
         s = upstream_states
-        assert s.skipped >= s.skipped_setup
-        assert s.success >= s.success_setup
+        assert s.skipped >= s.skipped_setup - indirect_upstream_setups
+        assert s.success >= s.success_setup - indirect_upstream_setups
         assert s.done == s.failed + s.success + s.removed + s.upstream_failed + s.skipped
 
+        setup_upstream_tasks = []
         with dag_maker() as dag:
             downstream = EmptyOperator(task_id="downstream", trigger_rule=trigger_rule)
             for i in range(5):
                 task = EmptyOperator(task_id=f"work_{i}", dag=dag)
                 task.set_downstream(downstream)
-            for i in range(upstream_setups):
-                task = EmptyOperator(task_id=f"setup_{i}", dag=dag).as_setup()
+            for i in range(direct_upstream_setups):
+                task = EmptyOperator(task_id=f"direct_setup_{i}", dag=dag).as_setup()
+                task.set_downstream(downstream)
+                setup_upstream_tasks.append(task)
+            for i in range(indirect_upstream_setups):
+                setup_task = EmptyOperator(task_id=f"indirect_setup_{i}", dag=dag).as_setup()
+                task = EmptyOperator(task_id=f"indirect_setup_downstream_{i}", dag=dag)
+                setup_task.set_downstream(task)
                 task.set_downstream(downstream)
+                setup_upstream_tasks.append(setup_task)
             assert task.start_date is not None
             run_date = task.start_date + datetime.timedelta(days=5)
 
@@ -1314,6 +1512,7 @@ def test_check_task_dependencies(
         dep_results = TriggerRuleDep()._evaluate_trigger_rule(
             ti=ti,
             dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
+            setup_upstream_tasks=setup_upstream_tasks,  # type: ignore
             session=dag_maker.session,
         )
         completed = all(dep.passed for dep in dep_results)
@@ -1331,37 +1530,37 @@ def test_check_task_dependencies(
             #
             # Tests for all_success
             #
-            ["all_success", _UpstreamTIStates(5, 0, 0, 0, 0, 0, 0, 0), True, None, True],
-            ["all_success", _UpstreamTIStates(2, 0, 0, 0, 0, 0, 0, 0), True, None, False],
-            ["all_success", _UpstreamTIStates(2, 0, 1, 0, 0, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["all_success", _UpstreamTIStates(2, 1, 0, 0, 0, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_success", _UpstreamTIStates(5, 0, 0, 0, 0, 0, 0, 0, 0), True, None, True],
+            ["all_success", _UpstreamTIStates(2, 0, 0, 0, 0, 0, 0, 0, 0), True, None, False],
+            ["all_success", _UpstreamTIStates(2, 0, 1, 0, 0, 0, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["all_success", _UpstreamTIStates(2, 1, 0, 0, 0, 0, 0, 0, 0), True, State.SKIPPED, False],
             # ti.map_index >= success
-            ["all_success", _UpstreamTIStates(3, 0, 0, 0, 2, 0, 0, 0), True, State.REMOVED, True],
+            ["all_success", _UpstreamTIStates(3, 0, 0, 0, 2, 0, 0, 0, 0), True, State.REMOVED, True],
             #
             # Tests for one_success
             #
-            ["one_success", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
-            ["one_success", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, True],
-            ["one_success", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, True],
-            ["one_success", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, True],
-            ["one_success", _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
-            ["one_success", _UpstreamTIStates(0, 4, 1, 0, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 3, 1, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 4, 0, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 0, 4, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 0, 0, 5, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["one_success", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, True],
+            ["one_success", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
+            ["one_success", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, True],
+            ["one_success", _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["one_success", _UpstreamTIStates(0, 4, 1, 0, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 3, 1, 1, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 4, 0, 1, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 0, 4, 1, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 0, 0, 5, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
             #
             # Tests for all_failed
             #
-            ["all_failed", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0), True, None, True],
-            ["all_failed", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_failed", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
             [
                 "all_failed",
-                _UpstreamTIStates(2, 1, 0, 0, 1, 4, 0, 0),
+                _UpstreamTIStates(2, 1, 0, 0, 1, 4, 0, 0, 0),
                 True,
                 State.SKIPPED,
                 False,
@@ -1369,14 +1568,14 @@ def test_check_task_dependencies(
             #
             # Tests for one_failed
             #
-            ["one_failed", _UpstreamTIStates(5, 0, 0, 0, 0, 0, 0, 0), True, None, False],
-            ["one_failed", _UpstreamTIStates(2, 0, 0, 0, 0, 0, 0, 0), True, None, False],
-            ["one_failed", _UpstreamTIStates(2, 0, 1, 0, 0, 0, 0, 0), True, None, True],
-            ["one_failed", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
-            ["one_failed", _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
+            ["one_failed", _UpstreamTIStates(5, 0, 0, 0, 0, 0, 0, 0, 0), True, None, False],
+            ["one_failed", _UpstreamTIStates(2, 0, 0, 0, 0, 0, 0, 0, 0), True, None, False],
+            ["one_failed", _UpstreamTIStates(2, 0, 1, 0, 0, 0, 0, 0, 0), True, None, True],
+            ["one_failed", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["one_failed", _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
             [
                 "one_failed",
-                _UpstreamTIStates(2, 2, 0, 0, 1, 5, 0, 0),
+                _UpstreamTIStates(2, 2, 0, 0, 1, 5, 0, 0, 0),
                 True,
                 State.SKIPPED,
                 False,
@@ -1384,10 +1583,10 @@ def test_check_task_dependencies(
             #
             # Tests for done
             #
-            ["all_done", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
-            ["all_done", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
-            ["all_done", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, False],
-            ["all_done", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
+            ["all_done", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_done", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
+            ["all_done", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
         ],
     )
     def test_check_task_dependencies_for_mapped(
@@ -1431,6 +1630,7 @@ def do_something_else(i):
         dep_results = TriggerRuleDep()._evaluate_trigger_rule(
             ti=ti,
             dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
+            setup_upstream_tasks=[],
             session=session,
         )
         completed = all(dep.passed for dep in dep_results)
diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py b/tests/ti_deps/deps/test_trigger_rule_dep.py
index faa70b5a4951d..6904c4161d37a 100644
--- a/tests/ti_deps/deps/test_trigger_rule_dep.py
+++ b/tests/ti_deps/deps/test_trigger_rule_dep.py
@@ -52,6 +52,7 @@ def _get_task_instance(
         done: int = 0,
         skipped_setup: int = 0,
         success_setup: int = 0,
+        failed_setup: int = 0,
         normal_tasks: list[str] | None = None,
         setup_tasks: list[str] | None = None,
     ):
@@ -78,6 +79,7 @@ def _get_task_instance(
             done=done,
             skipped_setup=skipped_setup,
             success_setup=success_setup,
+            failed_setup=failed_setup,
         )
         monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: fake_upstream_states)
 
@@ -158,6 +160,7 @@ def test_one_success_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -180,6 +183,7 @@ def test_one_success_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -203,6 +207,7 @@ def test_one_failure_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -226,6 +231,7 @@ def test_one_failure_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -248,6 +254,7 @@ def test_one_failure_tr_success_no_failed(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -270,6 +277,7 @@ def test_one_done_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -292,6 +300,7 @@ def test_one_done_tr_success_with_failed(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -314,6 +323,7 @@ def test_one_done_tr_skip(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -337,6 +347,7 @@ def test_one_done_tr_upstream_failed(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -361,6 +372,7 @@ def test_all_success_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -384,6 +396,7 @@ def test_all_success_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -412,6 +425,7 @@ def test_all_success_tr_skip(self, session, get_task_instance, flag_upstream_fai
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -444,6 +458,7 @@ def test_all_success_tr_skip_wait_for_past_depends_before_skipping(self, session
                     dep_context=DepContext(
                         flag_upstream_failed=True, wait_for_past_depends_before_skipping=True
                     ),
+                    setup_upstream_tasks=[],
                     session=session,
                 )
             )
@@ -478,6 +493,7 @@ def test_all_success_tr_skip_wait_for_past_depends_before_skipping_past_depends_
                     dep_context=DepContext(
                         flag_upstream_failed=True, wait_for_past_depends_before_skipping=True
                     ),
+                    setup_upstream_tasks=[],
                     session=session,
                 )
             )
@@ -504,6 +520,7 @@ def test_none_failed_tr_success(self, session, get_task_instance, flag_upstream_
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -528,6 +545,7 @@ def test_none_failed_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -552,6 +570,7 @@ def test_none_failed_min_one_success_tr_success(self, session, get_task_instance
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -575,6 +594,7 @@ def test_none_failed_min_one_success_tr_skipped(self, session, get_task_instance
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=True),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -599,6 +619,7 @@ def test_none_failed_min_one_success_tr_failure(self, session, get_task_instance
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -623,6 +644,7 @@ def test_all_failed_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -646,6 +668,7 @@ def test_all_failed_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -670,6 +693,7 @@ def test_all_done_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -694,14 +718,14 @@ def test_all_done_tr_success(self, session, get_task_instance):
             ),
             param(
                 dict(work=2, setup=1),
-                dict(success=2, done=3),
+                dict(success=2, done=3, failed_setup=1),
                 "requires at least one upstream setup task be successful",
                 UPSTREAM_FAILED,
                 id="setup failed",
             ),
             param(
                 dict(work=2, setup=2),
-                dict(success=2, done=4, success_setup=1),
+                dict(success=2, done=4, success_setup=1, failed_setup=1),
                 None,
                 None,
                 id="one setup failed one success",
@@ -715,14 +739,14 @@ def test_all_done_tr_success(self, session, get_task_instance):
             ),
             param(
                 dict(work=2, setup=1),
-                dict(success=2, done=3, failed=1),
+                dict(success=2, done=3, failed=1, failed_setup=1),
                 "requires at least one upstream setup task be successful",
                 UPSTREAM_FAILED,
                 id="setup failed",
             ),
             param(
                 dict(work=2, setup=2),
-                dict(success=2, done=4, failed=1, skipped_setup=1),
+                dict(success=2, done=4, failed=1, skipped_setup=1, failed_setup=1),
                 "requires at least one upstream setup task be successful",
                 UPSTREAM_FAILED,
                 id="one setup failed one skipped",
@@ -763,7 +787,7 @@ def test_teardown_tr_not_all_done(
         """
         All-done trigger rule success
         """
-        ti = get_task_instance(
+        ti: TaskInstance = get_task_instance(
             TriggerRule.ALL_DONE_SETUP_SUCCESS,
             **states,
             normal_tasks=[f"w{x}" for x in range(task_cfg["work"])],
@@ -771,7 +795,10 @@ def test_teardown_tr_not_all_done(
         )
         dep_statuses = tuple(
             TriggerRuleDep()._evaluate_trigger_rule(
-                ti=ti, dep_context=DepContext(flag_upstream_failed=True), session=session
+                ti=ti,
+                dep_context=DepContext(flag_upstream_failed=True),
+                setup_upstream_tasks=[t for t in ti.task.upstream_list if t.is_setup],
+                session=session,
             )
         )
         if exp_reason:
@@ -802,6 +829,7 @@ def test_all_skipped_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -827,6 +855,7 @@ def test_all_skipped_tr_success(self, session, get_task_instance, flag_upstream_
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -852,6 +881,7 @@ def test_all_done_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -877,6 +907,7 @@ def test_none_skipped_tr_success(self, session, get_task_instance, flag_upstream
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -901,6 +932,7 @@ def test_none_skipped_tr_failure(self, session, get_task_instance, flag_upstream
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -926,6 +958,7 @@ def test_none_skipped_tr_failure_empty(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -951,6 +984,7 @@ def test_unknown_tr(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -984,9 +1018,9 @@ def _get_finished_tis(task_id: str) -> Iterator[TaskInstance]:
             return (ti for ti in tis.values() if ti.task_id in tis[task_id].task.upstream_task_ids)
 
         # check handling with cases that tasks are triggered from backfill with no finished tasks
-        assert _UpstreamTIStates.calculate(_get_finished_tis("op2")) == (1, 0, 0, 0, 0, 1, 0, 0)
-        assert _UpstreamTIStates.calculate(_get_finished_tis("op4")) == (1, 0, 1, 0, 0, 2, 0, 0)
-        assert _UpstreamTIStates.calculate(_get_finished_tis("op5")) == (2, 0, 1, 0, 0, 3, 0, 0)
+        assert _UpstreamTIStates.calculate(_get_finished_tis("op2")) == (1, 0, 0, 0, 0, 1, 0, 0, 0)
+        assert _UpstreamTIStates.calculate(_get_finished_tis("op4")) == (1, 0, 1, 0, 0, 2, 0, 0, 0)
+        assert _UpstreamTIStates.calculate(_get_finished_tis("op5")) == (2, 0, 1, 0, 0, 3, 0, 0, 0)
 
         dr.update_state(session=session)
         assert dr.state == DagRunState.SUCCESS
@@ -1015,6 +1049,7 @@ def test_mapped_task_upstream_removed_with_all_success_trigger_rules(
             done=5,
             skipped_setup=0,
             success_setup=0,
+            failed_setup=0,
         )
         monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states)
 
@@ -1023,6 +1058,7 @@ def test_mapped_task_upstream_removed_with_all_success_trigger_rules(
                 ti=ti,
                 # Marks the task as removed if upstream is removed.
                 dep_context=DepContext(flag_upstream_failed=True),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -1054,6 +1090,7 @@ def test_mapped_task_upstream_removed_with_all_failed_trigger_rules(
             done=5,
             skipped_setup=0,
             success_setup=0,
+            failed_setup=0,
         )
         monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states)
 
@@ -1061,6 +1098,7 @@ def test_mapped_task_upstream_removed_with_all_failed_trigger_rules(
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -1096,6 +1134,7 @@ def test_mapped_task_upstream_removed_with_none_failed_trigger_rules(
             done=5,
             skipped_setup=0,
             success_setup=0,
+            failed_setup=0,
         )
         monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states)
 
@@ -1103,6 +1142,7 @@ def test_mapped_task_upstream_removed_with_none_failed_trigger_rules(
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
+                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -1210,6 +1250,7 @@ def tg(a):
         # of expansion does not fail the dependency-checking logic.
         ti=next(ti for ti in dr.task_instances if ti.task_id == "tg.t3" and ti.map_index == -1),
         dep_context=DepContext(),
+        setup_upstream_tasks=[],
         session=session,
     )
     results = list(result_iterator)

From 7ea522c181a551c7cbc0a9c3c6f0d0b2334f8272 Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Sat, 26 Aug 2023 01:25:52 +0200
Subject: [PATCH 02/22] - Fix a bug with mapped tasks - update the method which
 find all upstream setup tasks - update some code according to code review

---
 airflow/models/abstractoperator.py          |  5 +++--
 airflow/models/dag.py                       | 11 +++++++++++
 airflow/ti_deps/deps/trigger_rule_dep.py    | 20 +++++++++++++-------
 tests/ti_deps/deps/test_trigger_rule_dep.py |  1 +
 4 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
index d3833e8f9b119..81e6cd48eec27 100644
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -300,8 +300,9 @@ def get_upstreams_only_setups(self) -> Iterable[Operator]:
         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.
         """
-        upstream_setups = {x for x in self.get_flat_relatives(upstream=True) if x.is_setup}
-        yield from upstream_setups
+        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.
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index 75fee04145de7..8cb2c8ffc94cb 100644
--- a/airflow/models/dag.py
+++ b/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,16 @@ 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 not in [
+                        TriggerRule.ALL_SUCCESS,
+                        TriggerRule.ONE_SUCCESS,
+                    ]:
+                        # this is required to ensure consistent clearing behavior when upstream
+                        raise ValueError(
+                            "Setup tasks must be followed with trigger rule ALL_SUCCESS or ONE_SUCCESS."
+                        )
             FailStopDagInvalidTriggerRule.check(dag=self, trigger_rule=task.trigger_rule)
 
     def __repr__(self):
diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index 4a92dc2059c5d..b84f9d072dd45 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -167,7 +167,14 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool:
             """Whether a task instance is a "relevant upstream" of the current task."""
             # All the setup tasks upstreams are relevant event if they are not a direct upstream.
             if upstream.task_id in map(lambda t: t.task_id, setup_upstream_tasks):
-                return True
+                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
             # Not actually an upstream task.
             if upstream.task_id not in task.upstream_task_ids:
                 return False
@@ -197,8 +204,6 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool:
         )
         upstream_states = _UpstreamTIStates.calculate(finished_upstream_tis)
 
-        upstream_setup = len(setup_upstream_tasks)  # count of setup tasks upstream of this task
-
         success = upstream_states.success
         skipped = upstream_states.skipped
         failed = upstream_states.failed
@@ -242,6 +247,7 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
         # "simple" tasks (no task or task group mapping involved).
         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)
         else:
             task_id_counts = session.execute(
                 select(TaskInstance.task_id, func.count(TaskInstance.task_id))
@@ -250,15 +256,15 @@ 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)
 
         upstream_done = done >= upstream
         setup_done = (success_setup + skipped_setup + failed_setup) >= upstream_setup
-        is_tear_down = task.is_teardown or trigger_rule == TR.ALL_DONE_SETUP_SUCCESS
 
         changed = False
         new_state = None
         if dep_context.flag_upstream_failed:
-            if not is_tear_down and failed_setup:
+            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
@@ -308,7 +314,7 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
                     # 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 is_tear_down and upstream_setup and not new_state and setup_done and skipped_setup > 0:
+            if not task.is_teardown and upstream_setup and not new_state 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
 
@@ -330,7 +336,7 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
         # 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 is_tear_down and (success_setup + skipped_setup) < upstream_setup:
+        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."
diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py b/tests/ti_deps/deps/test_trigger_rule_dep.py
index 6904c4161d37a..831dcafda82cb 100644
--- a/tests/ti_deps/deps/test_trigger_rule_dep.py
+++ b/tests/ti_deps/deps/test_trigger_rule_dep.py
@@ -793,6 +793,7 @@ def test_teardown_tr_not_all_done(
             normal_tasks=[f"w{x}" for x in range(task_cfg["work"])],
             setup_tasks=[f"s{x}" for x in range(task_cfg["setup"])],
         )
+        ti.task.is_teardown = True
         dep_statuses = tuple(
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,

From aed786abd92bded537502c78fe364259c99be042 Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Sat, 26 Aug 2023 02:37:24 +0200
Subject: [PATCH 03/22] Fix a bug and failed tests

---
 airflow/ti_deps/deps/trigger_rule_dep.py | 8 ++++++--
 tests/models/test_taskinstance.py        | 2 ++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index dd5c93deec2e6..e7c725923692f 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -245,9 +245,13 @@ 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 = len(setup_upstream_tasks)  # count of setup tasks upstream of this task
+        else:
+            upstream_setup = None
         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)
+            upstream_setup = upstream_setup or 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))
@@ -256,7 +260,7 @@ 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)
+            upstream_setup = upstream_setup or 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
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 9552232895b78..33e917ca6e5e4 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1490,6 +1490,8 @@ def test_check_task_dependencies(
         setup_upstream_tasks = []
         with dag_maker() as dag:
             downstream = EmptyOperator(task_id="downstream", trigger_rule=trigger_rule)
+            if trigger_rule == "all_done_setup_success":
+                downstream.is_teardown = True
             for i in range(5):
                 task = EmptyOperator(task_id=f"work_{i}", dag=dag)
                 task.set_downstream(downstream)

From 5ea28b62e20e2b60021115776e1b0978e1c24d06 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Sat, 26 Aug 2023 08:04:51 -0700
Subject: [PATCH 04/22] reduce diffs by making setup_upstream_tasks optional

---
 airflow/ti_deps/deps/trigger_rule_dep.py    |  5 +++-
 tests/ti_deps/deps/test_trigger_rule_dep.py | 33 ---------------------
 2 files changed, 4 insertions(+), 34 deletions(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index e7c725923692f..b9c5f2262414f 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -117,7 +117,7 @@ def _evaluate_trigger_rule(
         *,
         ti: TaskInstance,
         dep_context: DepContext,
-        setup_upstream_tasks: list[Operator],
+        setup_upstream_tasks: list[Operator] | None = None,
         session: Session,
     ) -> Iterator[TIDepStatus]:
         """Evaluate whether ``ti``'s trigger rule was met.
@@ -132,6 +132,7 @@ def _evaluate_trigger_rule(
         from airflow.models.operator import needs_expansion
         from airflow.models.taskinstance import TaskInstance
 
+        setup_upstream_tasks = setup_upstream_tasks or []
         task = ti.task
         upstream_tasks = {t.task_id: t for t in task.upstream_list}
         trigger_rule = task.trigger_rule
@@ -166,6 +167,8 @@ def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None:
         def _is_relevant_upstream(upstream: TaskInstance) -> bool:
             """Whether a task instance is a "relevant upstream" of the current task."""
             # All the setup tasks upstreams are relevant event if they are not a direct upstream.
+            if TYPE_CHECKING:
+                assert isinstance(setup_upstream_tasks, list)
             if upstream.task_id in map(lambda t: t.task_id, setup_upstream_tasks):
                 relevant = _get_relevant_upstream_map_indexes(upstream.task_id)
                 if relevant is None:
diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py b/tests/ti_deps/deps/test_trigger_rule_dep.py
index 831dcafda82cb..fe17fb750e999 100644
--- a/tests/ti_deps/deps/test_trigger_rule_dep.py
+++ b/tests/ti_deps/deps/test_trigger_rule_dep.py
@@ -160,7 +160,6 @@ def test_one_success_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -183,7 +182,6 @@ def test_one_success_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -207,7 +205,6 @@ def test_one_failure_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -231,7 +228,6 @@ def test_one_failure_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -254,7 +250,6 @@ def test_one_failure_tr_success_no_failed(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -277,7 +272,6 @@ def test_one_done_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -300,7 +294,6 @@ def test_one_done_tr_success_with_failed(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -323,7 +316,6 @@ def test_one_done_tr_skip(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -347,7 +339,6 @@ def test_one_done_tr_upstream_failed(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -372,7 +363,6 @@ def test_all_success_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -396,7 +386,6 @@ def test_all_success_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -425,7 +414,6 @@ def test_all_success_tr_skip(self, session, get_task_instance, flag_upstream_fai
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -458,7 +446,6 @@ def test_all_success_tr_skip_wait_for_past_depends_before_skipping(self, session
                     dep_context=DepContext(
                         flag_upstream_failed=True, wait_for_past_depends_before_skipping=True
                     ),
-                    setup_upstream_tasks=[],
                     session=session,
                 )
             )
@@ -493,7 +480,6 @@ def test_all_success_tr_skip_wait_for_past_depends_before_skipping_past_depends_
                     dep_context=DepContext(
                         flag_upstream_failed=True, wait_for_past_depends_before_skipping=True
                     ),
-                    setup_upstream_tasks=[],
                     session=session,
                 )
             )
@@ -520,7 +506,6 @@ def test_none_failed_tr_success(self, session, get_task_instance, flag_upstream_
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -545,7 +530,6 @@ def test_none_failed_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -570,7 +554,6 @@ def test_none_failed_min_one_success_tr_success(self, session, get_task_instance
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -594,7 +577,6 @@ def test_none_failed_min_one_success_tr_skipped(self, session, get_task_instance
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=True),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -619,7 +601,6 @@ def test_none_failed_min_one_success_tr_failure(self, session, get_task_instance
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -644,7 +625,6 @@ def test_all_failed_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -668,7 +648,6 @@ def test_all_failed_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -693,7 +672,6 @@ def test_all_done_tr_success(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -830,7 +808,6 @@ def test_all_skipped_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -856,7 +833,6 @@ def test_all_skipped_tr_success(self, session, get_task_instance, flag_upstream_
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -882,7 +858,6 @@ def test_all_done_tr_failure(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -908,7 +883,6 @@ def test_none_skipped_tr_success(self, session, get_task_instance, flag_upstream
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -933,7 +907,6 @@ def test_none_skipped_tr_failure(self, session, get_task_instance, flag_upstream
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -959,7 +932,6 @@ def test_none_skipped_tr_failure_empty(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -985,7 +957,6 @@ def test_unknown_tr(self, session, get_task_instance):
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -1059,7 +1030,6 @@ def test_mapped_task_upstream_removed_with_all_success_trigger_rules(
                 ti=ti,
                 # Marks the task as removed if upstream is removed.
                 dep_context=DepContext(flag_upstream_failed=True),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -1099,7 +1069,6 @@ def test_mapped_task_upstream_removed_with_all_failed_trigger_rules(
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -1143,7 +1112,6 @@ def test_mapped_task_upstream_removed_with_none_failed_trigger_rules(
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=False),
-                setup_upstream_tasks=[],
                 session=session,
             )
         )
@@ -1251,7 +1219,6 @@ def tg(a):
         # of expansion does not fail the dependency-checking logic.
         ti=next(ti for ti in dr.task_instances if ti.task_id == "tg.t3" and ti.map_index == -1),
         dep_context=DepContext(),
-        setup_upstream_tasks=[],
         session=session,
     )
     results = list(result_iterator)

From 4d7059cb12ea812452855cd491ba6a421f7e96c6 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Sat, 26 Aug 2023 08:24:49 -0700
Subject: [PATCH 05/22] if a task has no upstreams, there cannot be an upstream
 setup task, even indirect

---
 airflow/ti_deps/deps/trigger_rule_dep.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index b9c5f2262414f..998d3ed9d1e99 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -101,7 +101,7 @@ def _get_dep_statuses(
         # get all the setup tasks upstream of this task
         setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
         # Checking that all upstream dependencies have succeeded.
-        if not ti.task.upstream_task_ids and not setup_upstream_tasks:
+        if not ti.task.upstream_task_ids:
             yield self._passing_status(reason="The task instance did not have any upstream tasks.")
             return
         if ti.task.trigger_rule == TR.ALWAYS and not setup_upstream_tasks:

From b7a4c610b2177670c4fdc9349dae24dae1c447b7 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Sat, 26 Aug 2023 08:34:41 -0700
Subject: [PATCH 06/22] add indirect upstream tests as separate test

---
 tests/models/test_taskinstance.py | 181 ++++++++++++++++++++++++++++++
 1 file changed, 181 insertions(+)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 33e917ca6e5e4..c5a43143fd390 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1522,6 +1522,187 @@ def test_check_task_dependencies(
         assert completed == expect_passed
         assert ti.state == expect_state
 
+    @pytest.mark.parametrize(
+        "trigger_rule, direct_upstream_setups, indirect_upstream_setups, upstream_states,"
+        " flag_upstream_failed, expect_state, expect_passed",
+        [
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 2, 0, 0),
+                True,
+                None,
+                True,
+                id="indirect upstream setups - all success",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 0, 0),
+                True,
+                None,
+                False,
+                id="indirect upstream setups - one not done and one success",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 1, 0),
+                True,
+                None,
+                False,
+                id="indirect upstream setups - one not done and one skipped",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 1, 0),
+                True,
+                TaskInstanceState.SKIPPED,
+                True,
+                id="indirect upstream setups - all done: one skipped and one success",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 0, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - all done: one success and one failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 1, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - all done: one skipped and one failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 0, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - one not done and one failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(0, 0, 1, 0, 0, 1, 2, 0, 0),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - all setup success but upstream failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(5, 0, 0, 0, 0, 5, 2, 0, 0),
+                True,
+                None,
+                False,
+                id="indirect upstream setups - all setup success but upstream still running",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - one setup failed but upstream still running",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(4, 0, 1, 0, 0, 5, 0, 2, 0),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - all setup skipped but upstream failed",
+            ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(0, 7, 0, 0, 0, 7, 0, 0, 1),
+                True,
+                TaskInstanceState.UPSTREAM_FAILED,
+                False,
+                id="indirect upstream setups - one setup failed but upstream skipped",
+            ),
+        ],
+    )
+    def test_check_task_dependencies_indirect_upstream(
+        self,
+        monkeypatch,
+        dag_maker,
+        trigger_rule: str,
+        direct_upstream_setups: int,
+        indirect_upstream_setups: int,
+        upstream_states: _UpstreamTIStates,
+        flag_upstream_failed: bool,
+        expect_state: State,
+        expect_passed: bool,
+    ):
+        monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states)
+
+        # sanity checks
+        s = upstream_states
+        assert s.skipped >= s.skipped_setup - indirect_upstream_setups
+        assert s.success >= s.success_setup - indirect_upstream_setups
+        assert s.done == s.failed + s.success + s.removed + s.upstream_failed + s.skipped
+
+        setup_upstream_tasks = []
+        with dag_maker() as dag:
+            downstream = EmptyOperator(task_id="downstream", trigger_rule=trigger_rule)
+            if trigger_rule == "all_done_setup_success":
+                downstream.is_teardown = True
+            for i in range(5):
+                task = EmptyOperator(task_id=f"work_{i}", dag=dag)
+                task.set_downstream(downstream)
+            for i in range(direct_upstream_setups):
+                task = EmptyOperator(task_id=f"direct_setup_{i}", dag=dag).as_setup()
+                task.set_downstream(downstream)
+                setup_upstream_tasks.append(task)
+            for i in range(indirect_upstream_setups):
+                setup_task = EmptyOperator(task_id=f"indirect_setup_{i}", dag=dag).as_setup()
+                task = EmptyOperator(task_id=f"indirect_setup_downstream_{i}", dag=dag)
+                setup_task.set_downstream(task)
+                task.set_downstream(downstream)
+                setup_upstream_tasks.append(setup_task)
+            assert task.start_date is not None
+            run_date = task.start_date + datetime.timedelta(days=5)
+
+        ti = dag_maker.create_dagrun(execution_date=run_date).get_task_instance(downstream.task_id)
+        ti.task = downstream
+
+        dep_results = TriggerRuleDep()._evaluate_trigger_rule(
+            ti=ti,
+            dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
+            setup_upstream_tasks=setup_upstream_tasks,  # type: ignore
+            session=dag_maker.session,
+        )
+        completed = all(dep.passed for dep in dep_results)
+
+        assert completed == expect_passed
+        assert ti.state == expect_state
+
     # Parameterized tests to check for the correct firing
     # of the trigger_rule under various circumstances of mapped task
     # Numeric fields are in order:

From 45d5ae2ddef9c31d6e9c0beb9ce3df70e2a98132 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Sat, 26 Aug 2023 08:39:14 -0700
Subject: [PATCH 07/22] remove new tests from existing test

---
 tests/models/test_taskinstance.py | 158 +-----------------------------
 1 file changed, 5 insertions(+), 153 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index c5a43143fd390..4b571d10eeff0 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1345,126 +1345,6 @@ def test_depends_on_past(self, dag_maker):
                 False,
                 id="not all done, one skipped",
             ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 2, 0, 0),
-                True,
-                None,
-                True,
-                id="indirect upstream setups - all success",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 0, 0),
-                True,
-                None,
-                False,
-                id="indirect upstream setups - one not done and one success",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 1, 0),
-                True,
-                None,
-                False,
-                id="indirect upstream setups - one not done and one skipped",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 1, 0),
-                True,
-                TaskInstanceState.SKIPPED,
-                True,
-                id="indirect upstream setups - all done: one skipped and one success",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 0, 1),
-                True,
-                TaskInstanceState.UPSTREAM_FAILED,
-                False,
-                id="indirect upstream setups - all done: one success and one failed",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 1, 1),
-                True,
-                TaskInstanceState.UPSTREAM_FAILED,
-                False,
-                id="indirect upstream setups - all done: one skipped and one failed",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 0, 1),
-                True,
-                TaskInstanceState.UPSTREAM_FAILED,
-                False,
-                id="indirect upstream setups - one not done and one failed",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(0, 0, 1, 0, 0, 1, 2, 0, 0),
-                True,
-                TaskInstanceState.UPSTREAM_FAILED,
-                False,
-                id="indirect upstream setups - all setup success but upstream failed",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(5, 0, 0, 0, 0, 5, 2, 0, 0),
-                True,
-                None,
-                False,
-                id="indirect upstream setups - all setup success but upstream still running",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 1),
-                True,
-                TaskInstanceState.UPSTREAM_FAILED,
-                False,
-                id="indirect upstream setups - one setup failed but upstream still running",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(4, 0, 1, 0, 0, 5, 0, 2, 0),
-                True,
-                TaskInstanceState.UPSTREAM_FAILED,
-                False,
-                id="indirect upstream setups - all setup skipped but upstream failed",
-            ),
-            param(
-                "all_success",
-                0,
-                2,
-                _UpstreamTIStates(0, 7, 0, 0, 0, 7, 0, 0, 1),
-                True,
-                TaskInstanceState.UPSTREAM_FAILED,
-                False,
-                id="indirect upstream setups - one setup failed but upstream skipped",
-            ),
         ],
     )
     def test_check_task_dependencies(
@@ -1472,8 +1352,7 @@ def test_check_task_dependencies(
         monkeypatch,
         dag_maker,
         trigger_rule: str,
-        direct_upstream_setups: int,
-        indirect_upstream_setups: int,
+        upstream_setups: int,
         upstream_states: _UpstreamTIStates,
         flag_upstream_failed: bool,
         expect_state: State,
@@ -1483,44 +1362,17 @@ def test_check_task_dependencies(
 
         # sanity checks
         s = upstream_states
-        assert s.skipped >= s.skipped_setup - indirect_upstream_setups
-        assert s.success >= s.success_setup - indirect_upstream_setups
+        assert s.skipped >= s.skipped_setup
+        assert s.success >= s.success_setup
         assert s.done == s.failed + s.success + s.removed + s.upstream_failed + s.skipped
 
-        setup_upstream_tasks = []
         with dag_maker() as dag:
             downstream = EmptyOperator(task_id="downstream", trigger_rule=trigger_rule)
-            if trigger_rule == "all_done_setup_success":
-                downstream.is_teardown = True
             for i in range(5):
                 task = EmptyOperator(task_id=f"work_{i}", dag=dag)
                 task.set_downstream(downstream)
-            for i in range(direct_upstream_setups):
-                task = EmptyOperator(task_id=f"direct_setup_{i}", dag=dag).as_setup()
-                task.set_downstream(downstream)
-                setup_upstream_tasks.append(task)
-            for i in range(indirect_upstream_setups):
-                setup_task = EmptyOperator(task_id=f"indirect_setup_{i}", dag=dag).as_setup()
-                task = EmptyOperator(task_id=f"indirect_setup_downstream_{i}", dag=dag)
-                setup_task.set_downstream(task)
-                task.set_downstream(downstream)
-                setup_upstream_tasks.append(setup_task)
-            assert task.start_date is not None
-            run_date = task.start_date + datetime.timedelta(days=5)
-
-        ti = dag_maker.create_dagrun(execution_date=run_date).get_task_instance(downstream.task_id)
-        ti.task = downstream
-
-        dep_results = TriggerRuleDep()._evaluate_trigger_rule(
-            ti=ti,
-            dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
-            setup_upstream_tasks=setup_upstream_tasks,  # type: ignore
-            session=dag_maker.session,
-        )
-        completed = all(dep.passed for dep in dep_results)
-
-        assert completed == expect_passed
-        assert ti.state == expect_state
+            for i in range(upstream_setups):
+                task = EmptyOperator(task_id=f"setup_{i}", dag=dag).as_setup()
 
     @pytest.mark.parametrize(
         "trigger_rule, direct_upstream_setups, indirect_upstream_setups, upstream_states,"

From 8308a4054f1542d12a92b626eb1645b30ee770ff Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Sat, 26 Aug 2023 09:07:08 -0700
Subject: [PATCH 08/22] remove setup_upstream_tasks from evaluate_trigger_rule
 and restore TR.ALWAYS behavior

---
 airflow/ti_deps/deps/trigger_rule_dep.py    | 17 +++++------------
 tests/models/test_taskinstance.py           |  5 -----
 tests/ti_deps/deps/test_trigger_rule_dep.py |  1 -
 3 files changed, 5 insertions(+), 18 deletions(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index 998d3ed9d1e99..6b80d5444526a 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -34,7 +34,6 @@
     from sqlalchemy.orm import Session
     from sqlalchemy.sql.expression import ColumnOperators
 
-    from airflow.models import Operator
     from airflow.models.taskinstance import TaskInstance
 
 
@@ -98,33 +97,27 @@ def _get_dep_statuses(
         session: Session,
         dep_context: DepContext,
     ) -> Iterator[TIDepStatus]:
-        # get all the setup tasks upstream of this task
-        setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
         # Checking that all upstream dependencies have succeeded.
         if not ti.task.upstream_task_ids:
             yield self._passing_status(reason="The task instance did not have any upstream tasks.")
             return
-        if ti.task.trigger_rule == TR.ALWAYS and not setup_upstream_tasks:
+        if ti.task.trigger_rule == TR.ALWAYS:
             # even with ALWAYS trigger rule, we still need to check setup tasks
             yield self._passing_status(reason="The task had a always trigger rule set.")
             return
-        yield from self._evaluate_trigger_rule(
-            ti=ti, dep_context=dep_context, setup_upstream_tasks=setup_upstream_tasks, session=session
-        )
+        yield from self._evaluate_trigger_rule(ti=ti, dep_context=dep_context, session=session)
 
     def _evaluate_trigger_rule(
         self,
         *,
         ti: TaskInstance,
         dep_context: DepContext,
-        setup_upstream_tasks: list[Operator] | None = None,
         session: Session,
     ) -> Iterator[TIDepStatus]:
         """Evaluate whether ``ti``'s trigger rule was met.
 
         :param ti: Task instance to evaluate the trigger rule of.
         :param dep_context: The current dependency context.
-        setup_upstream_tasks: The setup tasks upstream of the current task.
         :param session: Database session.
         """
         from airflow.models.abstractoperator import NotMapped
@@ -132,7 +125,9 @@ def _evaluate_trigger_rule(
         from airflow.models.operator import needs_expansion
         from airflow.models.taskinstance import TaskInstance
 
-        setup_upstream_tasks = setup_upstream_tasks or []
+        # get all the setup tasks upstream of this task
+        setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
+
         task = ti.task
         upstream_tasks = {t.task_id: t for t in task.upstream_list}
         trigger_rule = task.trigger_rule
@@ -167,8 +162,6 @@ def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None:
         def _is_relevant_upstream(upstream: TaskInstance) -> bool:
             """Whether a task instance is a "relevant upstream" of the current task."""
             # All the setup tasks upstreams are relevant event if they are not a direct upstream.
-            if TYPE_CHECKING:
-                assert isinstance(setup_upstream_tasks, list)
             if upstream.task_id in map(lambda t: t.task_id, setup_upstream_tasks):
                 relevant = _get_relevant_upstream_map_indexes(upstream.task_id)
                 if relevant is None:
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 4b571d10eeff0..e6e0f65f2b07b 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1520,7 +1520,6 @@ def test_check_task_dependencies_indirect_upstream(
         assert s.success >= s.success_setup - indirect_upstream_setups
         assert s.done == s.failed + s.success + s.removed + s.upstream_failed + s.skipped
 
-        setup_upstream_tasks = []
         with dag_maker() as dag:
             downstream = EmptyOperator(task_id="downstream", trigger_rule=trigger_rule)
             if trigger_rule == "all_done_setup_success":
@@ -1531,13 +1530,11 @@ def test_check_task_dependencies_indirect_upstream(
             for i in range(direct_upstream_setups):
                 task = EmptyOperator(task_id=f"direct_setup_{i}", dag=dag).as_setup()
                 task.set_downstream(downstream)
-                setup_upstream_tasks.append(task)
             for i in range(indirect_upstream_setups):
                 setup_task = EmptyOperator(task_id=f"indirect_setup_{i}", dag=dag).as_setup()
                 task = EmptyOperator(task_id=f"indirect_setup_downstream_{i}", dag=dag)
                 setup_task.set_downstream(task)
                 task.set_downstream(downstream)
-                setup_upstream_tasks.append(setup_task)
             assert task.start_date is not None
             run_date = task.start_date + datetime.timedelta(days=5)
 
@@ -1547,7 +1544,6 @@ def test_check_task_dependencies_indirect_upstream(
         dep_results = TriggerRuleDep()._evaluate_trigger_rule(
             ti=ti,
             dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
-            setup_upstream_tasks=setup_upstream_tasks,  # type: ignore
             session=dag_maker.session,
         )
         completed = all(dep.passed for dep in dep_results)
@@ -1665,7 +1661,6 @@ def do_something_else(i):
         dep_results = TriggerRuleDep()._evaluate_trigger_rule(
             ti=ti,
             dep_context=DepContext(flag_upstream_failed=flag_upstream_failed),
-            setup_upstream_tasks=[],
             session=session,
         )
         completed = all(dep.passed for dep in dep_results)
diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py b/tests/ti_deps/deps/test_trigger_rule_dep.py
index fe17fb750e999..58c94dc348151 100644
--- a/tests/ti_deps/deps/test_trigger_rule_dep.py
+++ b/tests/ti_deps/deps/test_trigger_rule_dep.py
@@ -776,7 +776,6 @@ def test_teardown_tr_not_all_done(
             TriggerRuleDep()._evaluate_trigger_rule(
                 ti=ti,
                 dep_context=DepContext(flag_upstream_failed=True),
-                setup_upstream_tasks=[t for t in ti.task.upstream_list if t.is_setup],
                 session=session,
             )
         )

From 4101b256c474e5ce1766d8184453cda6309ddf2c Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Sat, 26 Aug 2023 09:14:18 -0700
Subject: [PATCH 09/22] remove comment

---
 airflow/ti_deps/deps/trigger_rule_dep.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index 6b80d5444526a..c824657729d38 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -125,7 +125,6 @@ def _evaluate_trigger_rule(
         from airflow.models.operator import needs_expansion
         from airflow.models.taskinstance import TaskInstance
 
-        # get all the setup tasks upstream of this task
         setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
 
         task = ti.task

From c85cd23f149e42407a24062e9208848331464eb0 Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Sat, 26 Aug 2023 23:23:49 +0200
Subject: [PATCH 10/22] fix trigger_rule_dep test

---
 tests/ti_deps/deps/test_trigger_rule_dep.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py b/tests/ti_deps/deps/test_trigger_rule_dep.py
index 58c94dc348151..99502c02aee0e 100644
--- a/tests/ti_deps/deps/test_trigger_rule_dep.py
+++ b/tests/ti_deps/deps/test_trigger_rule_dep.py
@@ -989,9 +989,9 @@ def _get_finished_tis(task_id: str) -> Iterator[TaskInstance]:
             return (ti for ti in tis.values() if ti.task_id in tis[task_id].task.upstream_task_ids)
 
         # check handling with cases that tasks are triggered from backfill with no finished tasks
-        assert _UpstreamTIStates.calculate(_get_finished_tis("op2")) == (1, 0, 0, 0, 0, 1, 0, 0, 0)
-        assert _UpstreamTIStates.calculate(_get_finished_tis("op4")) == (1, 0, 1, 0, 0, 2, 0, 0, 0)
-        assert _UpstreamTIStates.calculate(_get_finished_tis("op5")) == (2, 0, 1, 0, 0, 3, 0, 0, 0)
+        assert _UpstreamTIStates.calculate(_get_finished_tis("op2"), []) == (1, 0, 0, 0, 0, 1, 0, 0, 0)
+        assert _UpstreamTIStates.calculate(_get_finished_tis("op4"), []) == (1, 0, 1, 0, 0, 2, 0, 0, 0)
+        assert _UpstreamTIStates.calculate(_get_finished_tis("op5"), []) == (2, 0, 1, 0, 0, 3, 0, 0, 0)
 
         dr.update_state(session=session)
         assert dr.state == DagRunState.SUCCESS

From f9262ecde71da9b2b496dc9f746e36bbe634779e Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Sat, 26 Aug 2023 23:24:44 +0200
Subject: [PATCH 11/22] Split upstream and setup upstream in two lists to fix a
 bug in trigger rule dep check

---
 airflow/ti_deps/deps/trigger_rule_dep.py | 46 ++++++++++-----
 tests/models/test_taskinstance.py        | 74 +++++++++---------------
 2 files changed, 58 insertions(+), 62 deletions(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index c824657729d38..1b3dedbaec7d3 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -54,7 +54,9 @@ class _UpstreamTIStates(NamedTuple):
     failed_setup: int
 
     @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
@@ -68,8 +70,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),
@@ -125,7 +128,10 @@ def _evaluate_trigger_rule(
         from airflow.models.operator import needs_expansion
         from airflow.models.taskinstance import TaskInstance
 
-        setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
+        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())
 
         task = ti.task
         upstream_tasks = {t.task_id: t for t in task.upstream_list}
@@ -160,16 +166,6 @@ def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None:
 
         def _is_relevant_upstream(upstream: TaskInstance) -> bool:
             """Whether a task instance is a "relevant upstream" of the current task."""
-            # All the setup tasks upstreams are relevant event if they are not a direct upstream.
-            if upstream.task_id in map(lambda t: t.task_id, setup_upstream_tasks):
-                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
             # Not actually an upstream task.
             if upstream.task_id not in task.upstream_task_ids:
                 return False
@@ -192,12 +188,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 setup" of the current task."""
+            if upstream.task_id in map(lambda t: t.task_id, setup_upstream_tasks):
+                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
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index e6e0f65f2b07b..16bcf727e8661 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1103,36 +1103,33 @@ def test_depends_on_past(self, dag_maker):
     # Numeric fields are in order:
     #   successes, skipped, failed, upstream_failed, removed, done
     @pytest.mark.parametrize(
-        "trigger_rule, direct_upstream_setups, indirect_upstream_setups, upstream_states,"
-        " flag_upstream_failed, expect_state, expect_passed",
+        "trigger_rule, upstream_setups,, upstream_states, flag_upstream_failed, expect_state, expect_passed",
         [
             #
             # Tests for all_success
             #
-            ["all_success", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_success", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
+            ["all_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
             [
                 "all_success",
                 0,
-                0,
                 _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0),
                 True,
                 State.UPSTREAM_FAILED,
                 False,
             ],
-            ["all_success", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
             #
             # Tests for one_success
             #
-            ["one_success", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["one_success", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, True],
-            ["one_success", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
-            ["one_success", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, True],
-            ["one_success", 0, 0, _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["one_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["one_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, True],
+            ["one_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
+            ["one_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, True],
+            ["one_success", 0, _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
             [
                 "one_success",
                 0,
-                0,
                 _UpstreamTIStates(0, 4, 1, 0, 0, 5, 0, 0, 0),
                 True,
                 State.UPSTREAM_FAILED,
@@ -1141,7 +1138,6 @@ def test_depends_on_past(self, dag_maker):
             [
                 "one_success",
                 0,
-                0,
                 _UpstreamTIStates(0, 3, 1, 1, 0, 5, 0, 0, 0),
                 True,
                 State.UPSTREAM_FAILED,
@@ -1150,7 +1146,6 @@ def test_depends_on_past(self, dag_maker):
             [
                 "one_success",
                 0,
-                0,
                 _UpstreamTIStates(0, 4, 0, 1, 0, 5, 0, 0, 0),
                 True,
                 State.UPSTREAM_FAILED,
@@ -1159,7 +1154,6 @@ def test_depends_on_past(self, dag_maker):
             [
                 "one_success",
                 0,
-                0,
                 _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0),
                 True,
                 State.UPSTREAM_FAILED,
@@ -1168,7 +1162,6 @@ def test_depends_on_past(self, dag_maker):
             [
                 "one_success",
                 0,
-                0,
                 _UpstreamTIStates(0, 0, 4, 1, 0, 5, 0, 0, 0),
                 True,
                 State.UPSTREAM_FAILED,
@@ -1177,7 +1170,6 @@ def test_depends_on_past(self, dag_maker):
             [
                 "one_success",
                 0,
-                0,
                 _UpstreamTIStates(0, 0, 0, 5, 0, 5, 0, 0, 0),
                 True,
                 State.UPSTREAM_FAILED,
@@ -1186,33 +1178,33 @@ def test_depends_on_past(self, dag_maker):
             #
             # Tests for all_failed
             #
-            ["all_failed", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", 0, 0, _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_failed", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_failed", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
             #
             # Tests for one_failed
             #
-            ["one_failed", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
-            ["one_failed", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
-            ["one_failed", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
-            ["one_failed", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
-            ["one_failed", 0, 0, _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["one_failed", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["one_failed", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
+            ["one_failed", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
+            ["one_failed", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["one_failed", 0, _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
             #
             # Tests for done
             #
-            ["all_done", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_done", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
-            ["all_done", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
-            ["all_done", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_done", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
+            ["all_done", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
             #
             # Tests for all_done_setup_success: no upstream setups -> same as all_done
             #
-            ["all_done_setup_success", 0, 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_done_setup_success", 0, 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
-            ["all_done_setup_success", 0, 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
-            ["all_done_setup_success", 0, 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done_setup_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
+            ["all_done_setup_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
+            ["all_done_setup_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done_setup_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
             #
             # Tests for all_done_setup_success: with upstream setups -> different from all_done
             #
@@ -1228,7 +1220,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                0,
                 _UpstreamTIStates(6, 0, 0, 0, 0, 6, 1, 0, 0),
                 True,
                 None,
@@ -1238,7 +1229,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                0,
                 _UpstreamTIStates(7, 0, 0, 0, 0, 7, 2, 0, 0),
                 True,
                 None,
@@ -1248,7 +1238,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                0,
                 _UpstreamTIStates(5, 0, 1, 0, 0, 6, 0, 0, 1),
                 True,
                 State.UPSTREAM_FAILED,
@@ -1258,7 +1247,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                0,
                 _UpstreamTIStates(5, 0, 2, 0, 0, 7, 0, 0, 2),
                 True,
                 State.UPSTREAM_FAILED,
@@ -1268,7 +1256,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                0,
                 _UpstreamTIStates(5, 1, 0, 0, 0, 6, 0, 1, 0),
                 True,
                 State.SKIPPED,
@@ -1278,7 +1265,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                0,
                 _UpstreamTIStates(5, 2, 0, 0, 0, 7, 0, 2, 0),
                 True,
                 State.SKIPPED,
@@ -1288,7 +1274,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                0,
                 _UpstreamTIStates(5, 1, 1, 0, 0, 7, 0, 1, 1),
                 True,
                 State.UPSTREAM_FAILED,
@@ -1298,7 +1283,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                0,
                 _UpstreamTIStates(6, 0, 1, 0, 0, 7, 1, 0, 1),
                 True,
                 None,
@@ -1308,7 +1292,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                0,
                 _UpstreamTIStates(6, 1, 0, 0, 0, 7, 1, 1, 0),
                 True,
                 None,
@@ -1318,7 +1301,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                0,
                 _UpstreamTIStates(3, 0, 0, 0, 0, 3, 1, 0, 0),
                 True,
                 None,
@@ -1328,7 +1310,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                0,
                 _UpstreamTIStates(3, 0, 1, 0, 0, 4, 1, 0, 0),
                 True,
                 None,
@@ -1338,7 +1319,6 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                0,
                 _UpstreamTIStates(3, 1, 0, 0, 0, 4, 1, 0, 0),
                 True,
                 None,

From d01e8026e866f36ee614f23602b3293511ca221e Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Sun, 27 Aug 2023 13:42:08 +0200
Subject: [PATCH 12/22] Make direct setup task respect trigger rule

---
 airflow/models/abstractoperator.py       |  4 ++--
 airflow/ti_deps/deps/trigger_rule_dep.py | 19 +++++++++++++------
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
index 81e6cd48eec27..da022028ef640 100644
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -293,7 +293,7 @@ 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]:
+    def get_indirect_upstreams_only_setups(self) -> Iterable[Operator]:
         """
         Only upstream setups.
 
@@ -301,7 +301,7 @@ def get_upstreams_only_setups(self) -> Iterable[Operator]:
         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:
+            if task.is_setup and task.task_id not in self.upstream_task_ids:
                 yield task
 
     def _iter_all_mapped_downstreams(self) -> Iterator[MappedOperator | MappedTaskGroup]:
diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index 1b3dedbaec7d3..854170207505b 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -131,7 +131,7 @@ def _evaluate_trigger_rule(
         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 = list(ti.task.get_indirect_upstreams_only_setups())
 
         task = ti.task
         upstream_tasks = {t.task_id: t for t in task.upstream_list}
@@ -189,8 +189,15 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool:
             return False
 
         def _is_relevant_setup_upstream(upstream: TaskInstance) -> bool:
-            """Whether a task instance is a "relevant setup" of the current task."""
+            """Whether a task instance is a "relevant indirect upstream setup" of the current task.
+
+            For teardown tasks, all relevant setup task should be considered as upstream setup.
+            """
             if upstream.task_id in map(lambda t: t.task_id, setup_upstream_tasks):
+                if upstream.task_id in ti.task.upstream_task_ids and not ti.task.is_teardown:
+                    # We should treat direct upstream setup tasks as normal upstream tasks,
+                    # except for teardown tasks
+                    return False
                 relevant = _get_relevant_upstream_map_indexes(upstream.task_id)
                 if relevant is None:
                     return True
@@ -258,11 +265,10 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
         # "simple" tasks (no task or task group mapping involved).
         if not ti.task.is_teardown:
             upstream_setup = len(setup_upstream_tasks)  # count of setup tasks upstream of this task
-        else:
-            upstream_setup = None
         if not any(needs_expansion(t) for t in upstream_tasks.values()):
             upstream = len(upstream_tasks)
-            upstream_setup = upstream_setup or 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))
@@ -271,7 +277,8 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
                 .group_by(TaskInstance.task_id)
             ).all()
             upstream = sum(count for _, count in task_id_counts)
-            upstream_setup = upstream_setup or 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

From e8cddfb5e84ac96c94dbbc387112db8c519011d2 Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Sun, 27 Aug 2023 22:36:30 +0200
Subject: [PATCH 13/22] Check setup tasks before the direct upstream tasks

---
 airflow/models/abstractoperator.py       |   4 +-
 airflow/ti_deps/deps/trigger_rule_dep.py | 108 +++++++++++------------
 tests/models/test_taskinstance.py        |  12 ++-
 3 files changed, 65 insertions(+), 59 deletions(-)

diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
index da022028ef640..81e6cd48eec27 100644
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -293,7 +293,7 @@ def get_upstreams_only_setups_and_teardowns(self) -> Iterable[Operator]:
                     if t.is_teardown and not t == self:
                         yield t
 
-    def get_indirect_upstreams_only_setups(self) -> Iterable[Operator]:
+    def get_upstreams_only_setups(self) -> Iterable[Operator]:
         """
         Only upstream setups.
 
@@ -301,7 +301,7 @@ def get_indirect_upstreams_only_setups(self) -> Iterable[Operator]:
         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 and task.task_id not in self.upstream_task_ids:
+            if task.is_setup:
                 yield task
 
     def _iter_all_mapped_downstreams(self) -> Iterator[MappedOperator | MappedTaskGroup]:
diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index 854170207505b..c3f2ae027ca60 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -131,7 +131,7 @@ def _evaluate_trigger_rule(
         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_indirect_upstreams_only_setups())
+            setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
 
         task = ti.task
         upstream_tasks = {t.task_id: t for t in task.upstream_list}
@@ -189,15 +189,8 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool:
             return False
 
         def _is_relevant_setup_upstream(upstream: TaskInstance) -> bool:
-            """Whether a task instance is a "relevant indirect upstream setup" of the current task.
-
-            For teardown tasks, all relevant setup task should be considered as upstream setup.
-            """
+            """Whether a task instance is a "relevant upstream setup" of the current task."""
             if upstream.task_id in map(lambda t: t.task_id, setup_upstream_tasks):
-                if upstream.task_id in ti.task.upstream_task_ids and not ti.task.is_teardown:
-                    # We should treat direct upstream setup tasks as normal upstream tasks,
-                    # except for teardown tasks
-                    return False
                 relevant = _get_relevant_upstream_map_indexes(upstream.task_id)
                 if relevant is None:
                     return True
@@ -290,55 +283,58 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]:
                 # 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 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 not task.is_teardown and upstream_setup and not new_state and setup_done and skipped_setup > 0:
+            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:
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 16bcf727e8661..41b03a9d67db8 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1464,7 +1464,7 @@ def test_check_task_dependencies(
                 2,
                 _UpstreamTIStates(4, 0, 1, 0, 0, 5, 0, 2, 0),
                 True,
-                TaskInstanceState.UPSTREAM_FAILED,
+                TaskInstanceState.SKIPPED,
                 False,
                 id="indirect upstream setups - all setup skipped but upstream failed",
             ),
@@ -1478,6 +1478,16 @@ def test_check_task_dependencies(
                 False,
                 id="indirect upstream setups - one setup failed but upstream skipped",
             ),
+            param(
+                "all_success",
+                0,
+                2,
+                _UpstreamTIStates(4, 0, 1, 0, 0, 5, 0, 1, 0),
+                True,
+                None,
+                False,
+                id="indirect upstream setups - one setup still running but upstream failed",
+            ),
         ],
     )
     def test_check_task_dependencies_indirect_upstream(

From 32add9e271923e0a88a94b801e89a3b340cdc7ec Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Sun, 27 Aug 2023 14:13:30 -0700
Subject: [PATCH 14/22] revert some changes to tests for easier review

---
 airflow/ti_deps/deps/trigger_rule_dep.py    |   8 +-
 tests/models/test_taskinstance.py           | 197 ++++++++------------
 tests/ti_deps/deps/test_trigger_rule_dep.py |   4 +-
 3 files changed, 80 insertions(+), 129 deletions(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index c3f2ae027ca60..2deb74d4a3526 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -20,7 +20,8 @@
 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
 
@@ -37,7 +38,8 @@
     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,7 +53,7 @@ class _UpstreamTIStates(NamedTuple):
     done: int
     success_setup: int
     skipped_setup: int
-    failed_setup: int
+    failed_setup: int = 0
 
     @classmethod
     def calculate(
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 41b03a9d67db8..4b6fcf39d5937 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1103,108 +1103,59 @@ def test_depends_on_past(self, dag_maker):
     # Numeric fields are in order:
     #   successes, skipped, failed, upstream_failed, removed, done
     @pytest.mark.parametrize(
-        "trigger_rule, upstream_setups,, upstream_states, flag_upstream_failed, expect_state, expect_passed",
+        "trigger_rule, upstream_setups, upstream_states, flag_upstream_failed, expect_state, expect_passed",
         [
             #
             # Tests for all_success
             #
-            ["all_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
-            [
-                "all_success",
-                0,
-                _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0),
-                True,
-                State.UPSTREAM_FAILED,
-                False,
-            ],
-            ["all_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
+            ["all_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
+            ["all_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["all_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
             #
             # Tests for one_success
             #
-            ["one_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["one_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, True],
-            ["one_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
-            ["one_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, True],
-            ["one_success", 0, _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
-            [
-                "one_success",
-                0,
-                _UpstreamTIStates(0, 4, 1, 0, 0, 5, 0, 0, 0),
-                True,
-                State.UPSTREAM_FAILED,
-                False,
-            ],
-            [
-                "one_success",
-                0,
-                _UpstreamTIStates(0, 3, 1, 1, 0, 5, 0, 0, 0),
-                True,
-                State.UPSTREAM_FAILED,
-                False,
-            ],
-            [
-                "one_success",
-                0,
-                _UpstreamTIStates(0, 4, 0, 1, 0, 5, 0, 0, 0),
-                True,
-                State.UPSTREAM_FAILED,
-                False,
-            ],
-            [
-                "one_success",
-                0,
-                _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0),
-                True,
-                State.UPSTREAM_FAILED,
-                False,
-            ],
-            [
-                "one_success",
-                0,
-                _UpstreamTIStates(0, 0, 4, 1, 0, 5, 0, 0, 0),
-                True,
-                State.UPSTREAM_FAILED,
-                False,
-            ],
-            [
-                "one_success",
-                0,
-                _UpstreamTIStates(0, 0, 0, 5, 0, 5, 0, 0, 0),
-                True,
-                State.UPSTREAM_FAILED,
-                False,
-            ],
+            ["one_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
+            ["one_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, True],
+            ["one_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, True],
+            ["one_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, True],
+            ["one_success", 0, _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
+            ["one_success", 0, _UpstreamTIStates(0, 4, 1, 0, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", 0, _UpstreamTIStates(0, 3, 1, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", 0, _UpstreamTIStates(0, 4, 0, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", 0, _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", 0, _UpstreamTIStates(0, 0, 4, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", 0, _UpstreamTIStates(0, 0, 0, 5, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
             #
             # Tests for all_failed
             #
-            ["all_failed", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", 0, _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_failed", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0), True, None, True],
+            ["all_failed", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
             #
             # Tests for one_failed
             #
-            ["one_failed", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
-            ["one_failed", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
-            ["one_failed", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
-            ["one_failed", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
-            ["one_failed", 0, _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["one_failed", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
+            ["one_failed", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
+            ["one_failed", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, True],
+            ["one_failed", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
+            ["one_failed", 0, _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
             #
             # Tests for done
             #
-            ["all_done", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_done", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
-            ["all_done", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
-            ["all_done", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
+            ["all_done", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
+            ["all_done", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, False],
+            ["all_done", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
             #
             # Tests for all_done_setup_success: no upstream setups -> same as all_done
             #
-            ["all_done_setup_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_done_setup_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
-            ["all_done_setup_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
-            ["all_done_setup_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done_setup_success", 0, _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
+            ["all_done_setup_success", 0, _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
+            ["all_done_setup_success", 0, _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, False],
+            ["all_done_setup_success", 0, _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
             #
             # Tests for all_done_setup_success: with upstream setups -> different from all_done
             #
@@ -1220,7 +1171,7 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(6, 0, 0, 0, 0, 6, 1, 0, 0),
+                _UpstreamTIStates(6, 0, 0, 0, 0, 6, 1, 0),
                 True,
                 None,
                 True,
@@ -1229,7 +1180,7 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 2, 0, 0),
+                _UpstreamTIStates(7, 0, 0, 0, 0, 7, 2, 0),
                 True,
                 None,
                 True,
@@ -1256,7 +1207,7 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(5, 1, 0, 0, 0, 6, 0, 1, 0),
+                _UpstreamTIStates(5, 1, 0, 0, 0, 6, 0, 1),
                 True,
                 State.SKIPPED,
                 False,
@@ -1265,7 +1216,7 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                _UpstreamTIStates(5, 2, 0, 0, 0, 7, 0, 2, 0),
+                _UpstreamTIStates(5, 2, 0, 0, 0, 7, 0, 2),
                 True,
                 State.SKIPPED,
                 False,
@@ -1292,7 +1243,7 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 2,
-                _UpstreamTIStates(6, 1, 0, 0, 0, 7, 1, 1, 0),
+                _UpstreamTIStates(6, 1, 0, 0, 0, 7, 1, 1),
                 True,
                 None,
                 True,
@@ -1301,7 +1252,7 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(3, 0, 0, 0, 0, 3, 1, 0, 0),
+                _UpstreamTIStates(3, 0, 0, 0, 0, 3, 1, 0),
                 True,
                 None,
                 False,
@@ -1310,7 +1261,7 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(3, 0, 1, 0, 0, 4, 1, 0, 0),
+                _UpstreamTIStates(3, 0, 1, 0, 0, 4, 1, 0),
                 True,
                 None,
                 False,
@@ -1319,7 +1270,7 @@ def test_depends_on_past(self, dag_maker):
             param(
                 "all_done_setup_success",
                 1,
-                _UpstreamTIStates(3, 1, 0, 0, 0, 4, 1, 0, 0),
+                _UpstreamTIStates(3, 1, 0, 0, 0, 4, 1, 0),
                 True,
                 None,
                 False,
@@ -1551,37 +1502,37 @@ def test_check_task_dependencies_indirect_upstream(
             #
             # Tests for all_success
             #
-            ["all_success", _UpstreamTIStates(5, 0, 0, 0, 0, 0, 0, 0, 0), True, None, True],
-            ["all_success", _UpstreamTIStates(2, 0, 0, 0, 0, 0, 0, 0, 0), True, None, False],
-            ["all_success", _UpstreamTIStates(2, 0, 1, 0, 0, 0, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["all_success", _UpstreamTIStates(2, 1, 0, 0, 0, 0, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_success", _UpstreamTIStates(5, 0, 0, 0, 0, 0, 0, 0), True, None, True],
+            ["all_success", _UpstreamTIStates(2, 0, 0, 0, 0, 0, 0, 0), True, None, False],
+            ["all_success", _UpstreamTIStates(2, 0, 1, 0, 0, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["all_success", _UpstreamTIStates(2, 1, 0, 0, 0, 0, 0, 0), True, State.SKIPPED, False],
             # ti.map_index >= success
-            ["all_success", _UpstreamTIStates(3, 0, 0, 0, 2, 0, 0, 0, 0), True, State.REMOVED, True],
+            ["all_success", _UpstreamTIStates(3, 0, 0, 0, 2, 0, 0, 0), True, State.REMOVED, True],
             #
             # Tests for one_success
             #
-            ["one_success", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["one_success", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, True],
-            ["one_success", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, True],
-            ["one_success", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, True],
-            ["one_success", _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
-            ["one_success", _UpstreamTIStates(0, 4, 1, 0, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 3, 1, 1, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 4, 0, 1, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 0, 4, 1, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
-            ["one_success", _UpstreamTIStates(0, 0, 0, 5, 0, 5, 0, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
+            ["one_success", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, True],
+            ["one_success", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, True],
+            ["one_success", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, True],
+            ["one_success", _UpstreamTIStates(0, 5, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
+            ["one_success", _UpstreamTIStates(0, 4, 1, 0, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 3, 1, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 4, 0, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 0, 4, 1, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
+            ["one_success", _UpstreamTIStates(0, 0, 0, 5, 0, 5, 0, 0), True, State.UPSTREAM_FAILED, False],
             #
             # Tests for all_failed
             #
-            ["all_failed", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_failed", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
-            ["all_failed", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", _UpstreamTIStates(0, 0, 5, 0, 0, 5, 0, 0), True, None, True],
+            ["all_failed", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
+            ["all_failed", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, State.SKIPPED, False],
             [
                 "all_failed",
-                _UpstreamTIStates(2, 1, 0, 0, 1, 4, 0, 0, 0),
+                _UpstreamTIStates(2, 1, 0, 0, 1, 4, 0, 0),
                 True,
                 State.SKIPPED,
                 False,
@@ -1589,14 +1540,14 @@ def test_check_task_dependencies_indirect_upstream(
             #
             # Tests for one_failed
             #
-            ["one_failed", _UpstreamTIStates(5, 0, 0, 0, 0, 0, 0, 0, 0), True, None, False],
-            ["one_failed", _UpstreamTIStates(2, 0, 0, 0, 0, 0, 0, 0, 0), True, None, False],
-            ["one_failed", _UpstreamTIStates(2, 0, 1, 0, 0, 0, 0, 0, 0), True, None, True],
-            ["one_failed", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
-            ["one_failed", _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0, 0), True, State.SKIPPED, False],
+            ["one_failed", _UpstreamTIStates(5, 0, 0, 0, 0, 0, 0, 0), True, None, False],
+            ["one_failed", _UpstreamTIStates(2, 0, 0, 0, 0, 0, 0, 0), True, None, False],
+            ["one_failed", _UpstreamTIStates(2, 0, 1, 0, 0, 0, 0, 0), True, None, True],
+            ["one_failed", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
+            ["one_failed", _UpstreamTIStates(2, 3, 0, 0, 0, 5, 0, 0), True, State.SKIPPED, False],
             [
                 "one_failed",
-                _UpstreamTIStates(2, 2, 0, 0, 1, 5, 0, 0, 0),
+                _UpstreamTIStates(2, 2, 0, 0, 1, 5, 0, 0),
                 True,
                 State.SKIPPED,
                 False,
@@ -1604,10 +1555,10 @@ def test_check_task_dependencies_indirect_upstream(
             #
             # Tests for done
             #
-            ["all_done", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 0), True, None, True],
-            ["all_done", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0, 0), True, None, False],
-            ["all_done", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0, 0), True, None, False],
-            ["all_done", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0, 0), True, None, False],
+            ["all_done", _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0), True, None, True],
+            ["all_done", _UpstreamTIStates(2, 0, 0, 0, 0, 2, 0, 0), True, None, False],
+            ["all_done", _UpstreamTIStates(2, 0, 1, 0, 0, 3, 0, 0), True, None, False],
+            ["all_done", _UpstreamTIStates(2, 1, 0, 0, 0, 3, 0, 0), True, None, False],
         ],
     )
     def test_check_task_dependencies_for_mapped(
diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py b/tests/ti_deps/deps/test_trigger_rule_dep.py
index 99502c02aee0e..6db40229e3fad 100644
--- a/tests/ti_deps/deps/test_trigger_rule_dep.py
+++ b/tests/ti_deps/deps/test_trigger_rule_dep.py
@@ -774,9 +774,7 @@ def test_teardown_tr_not_all_done(
         ti.task.is_teardown = True
         dep_statuses = tuple(
             TriggerRuleDep()._evaluate_trigger_rule(
-                ti=ti,
-                dep_context=DepContext(flag_upstream_failed=True),
-                session=session,
+                ti=ti, dep_context=DepContext(flag_upstream_failed=True), session=session
             )
         )
         if exp_reason:

From 8e21572f305afa79a88edd1896c92b802f91bc07 Mon Sep 17 00:00:00 2001
From: Hussein Awala <houssein.awala.96@gmail.com>
Date: Sun, 27 Aug 2023 23:51:51 +0200
Subject: [PATCH 15/22] Apply suggestions from code review

Co-authored-by: Daniel Standish <15932138+dstandish@users.noreply.github.com>
---
 airflow/models/dag.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index d45d456c940e8..b1ff0a26d556f 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -722,11 +722,10 @@ def validate_setup_teardown(self):
                 for down_task in task.downstream_list:
                     if not down_task.is_teardown and down_task.trigger_rule not in [
                         TriggerRule.ALL_SUCCESS,
-                        TriggerRule.ONE_SUCCESS,
                     ]:
                         # this is required to ensure consistent clearing behavior when upstream
                         raise ValueError(
-                            "Setup tasks must be followed with trigger rule ALL_SUCCESS or ONE_SUCCESS."
+                            "Setup tasks must be followed with trigger rule ALL_SUCCESS."
                         )
             FailStopDagInvalidTriggerRule.check(dag=self, trigger_rule=task.trigger_rule)
 

From 96950932561f7802ef9170b0d2ff96dbfc1bcbb2 Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Mon, 28 Aug 2023 00:25:07 +0200
Subject: [PATCH 16/22] fix unit tests and static checks

---
 airflow/models/dag.py                       |  4 +---
 tests/ti_deps/deps/test_trigger_rule_dep.py | 12 +++++++++---
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index b1ff0a26d556f..d4265b3eb8ea3 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -724,9 +724,7 @@ def validate_setup_teardown(self):
                         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."
-                        )
+                        raise ValueError("Setup tasks must be followed with trigger rule ALL_SUCCESS.")
             FailStopDagInvalidTriggerRule.check(dag=self, trigger_rule=task.trigger_rule)
 
     def __repr__(self):
diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py b/tests/ti_deps/deps/test_trigger_rule_dep.py
index 6db40229e3fad..8e3a424050664 100644
--- a/tests/ti_deps/deps/test_trigger_rule_dep.py
+++ b/tests/ti_deps/deps/test_trigger_rule_dep.py
@@ -987,9 +987,15 @@ def _get_finished_tis(task_id: str) -> Iterator[TaskInstance]:
             return (ti for ti in tis.values() if ti.task_id in tis[task_id].task.upstream_task_ids)
 
         # check handling with cases that tasks are triggered from backfill with no finished tasks
-        assert _UpstreamTIStates.calculate(_get_finished_tis("op2"), []) == (1, 0, 0, 0, 0, 1, 0, 0, 0)
-        assert _UpstreamTIStates.calculate(_get_finished_tis("op4"), []) == (1, 0, 1, 0, 0, 2, 0, 0, 0)
-        assert _UpstreamTIStates.calculate(_get_finished_tis("op5"), []) == (2, 0, 1, 0, 0, 3, 0, 0, 0)
+        assert _UpstreamTIStates.calculate(_get_finished_tis("op2"), []) == _UpstreamTIStates(
+            1, 0, 0, 0, 0, 1, 0, 0, 0
+        )
+        assert _UpstreamTIStates.calculate(_get_finished_tis("op4"), []) == _UpstreamTIStates(
+            1, 0, 1, 0, 0, 2, 0, 0, 0
+        )
+        assert _UpstreamTIStates.calculate(_get_finished_tis("op5"), []) == _UpstreamTIStates(
+            2, 0, 1, 0, 0, 3, 0, 0, 0
+        )
 
         dr.update_state(session=session)
         assert dr.state == DagRunState.SUCCESS

From 2591dd3acf3a906a281362be078b285fef5c751a Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Sun, 27 Aug 2023 23:36:07 -0700
Subject: [PATCH 17/22] remove comment

---
 airflow/ti_deps/deps/trigger_rule_dep.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index 2deb74d4a3526..0daf9896c1aab 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -107,7 +107,6 @@ def _get_dep_statuses(
             yield self._passing_status(reason="The task instance did not have any upstream tasks.")
             return
         if ti.task.trigger_rule == TR.ALWAYS:
-            # even with ALWAYS trigger rule, we still need to check setup tasks
             yield self._passing_status(reason="The task had a always trigger rule set.")
             return
         yield from self._evaluate_trigger_rule(ti=ti, dep_context=dep_context, session=session)

From 999dd990aaedb2cfb185c73649ff263338837bb9 Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Mon, 28 Aug 2023 19:36:22 +0200
Subject: [PATCH 18/22] calculate number of setup tasks which we should wait
 for from TIs instead of tasks list

---
 airflow/ti_deps/deps/trigger_rule_dep.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index 0daf9896c1aab..c2f8a451f0922 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -134,6 +134,13 @@ def _evaluate_trigger_rule(
         else:
             setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
 
+        setup_upstream_tasks_ids = [task.task_id for task in setup_upstream_tasks]
+        setup_upstream_task_instances = [
+            t
+            for t in ti.get_dagrun(session).get_task_instances(session=session)
+            if t.task_id in setup_upstream_tasks_ids
+        ]
+
         task = ti.task
         upstream_tasks = {t.task_id: t for t in task.upstream_list}
         trigger_rule = task.trigger_rule
@@ -258,7 +265,7 @@ 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 = len(setup_upstream_tasks)  # count of setup tasks upstream of this task
+            upstream_setup = len(setup_upstream_task_instances)  # count of setup tasks upstream of this task
         if not any(needs_expansion(t) for t in upstream_tasks.values()):
             upstream = len(upstream_tasks)
             if ti.task.is_teardown:

From 9eddecb575570ff677ee73664cd743df28e56db7 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Tue, 29 Aug 2023 11:05:39 -0700
Subject: [PATCH 19/22] small simplification

---
 airflow/models/dag.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index d4265b3eb8ea3..db6723a97ef84 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -720,9 +720,7 @@ def validate_setup_teardown(self):
         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 not in [
-                        TriggerRule.ALL_SUCCESS,
-                    ]:
+                    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)

From 6cb5e5e503e5ee92a053f86f6165f1987c417466 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Tue, 29 Aug 2023 14:43:22 -0700
Subject: [PATCH 20/22] simplify by using existing structure and convert to set

---
 airflow/ti_deps/deps/trigger_rule_dep.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index c2f8a451f0922..4cec91b7b3eb4 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -134,12 +134,12 @@ def _evaluate_trigger_rule(
         else:
             setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
 
-        setup_upstream_tasks_ids = [task.task_id for task in setup_upstream_tasks]
-        setup_upstream_task_instances = [
+        setup_upstream_tasks_ids = {task.task_id for task in setup_upstream_tasks}
+        setup_upstream_task_instances = {
             t
             for t in ti.get_dagrun(session).get_task_instances(session=session)
             if t.task_id in setup_upstream_tasks_ids
-        ]
+        }
 
         task = ti.task
         upstream_tasks = {t.task_id: t for t in task.upstream_list}
@@ -198,7 +198,7 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool:
 
         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 map(lambda t: t.task_id, setup_upstream_tasks):
+            if upstream.task_id in setup_upstream_tasks_ids:
                 relevant = _get_relevant_upstream_map_indexes(upstream.task_id)
                 if relevant is None:
                     return True

From f6e3d56e5f06f77fec9c5306cbeffb65addbd979 Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Wed, 30 Aug 2023 00:13:18 +0200
Subject: [PATCH 21/22] Optimise the perf by getting the count of the relevant
 TI instead of getting all tasks

---
 airflow/ti_deps/deps/trigger_rule_dep.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index 4cec91b7b3eb4..39b71ef7d16b0 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -25,6 +25,7 @@
 
 from sqlalchemy import and_, func, or_, select
 
+from airflow.models import DagRun
 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
@@ -135,11 +136,6 @@ def _evaluate_trigger_rule(
             setup_upstream_tasks = list(ti.task.get_upstreams_only_setups())
 
         setup_upstream_tasks_ids = {task.task_id for task in setup_upstream_tasks}
-        setup_upstream_task_instances = {
-            t
-            for t in ti.get_dagrun(session).get_task_instances(session=session)
-            if t.task_id in setup_upstream_tasks_ids
-        }
 
         task = ti.task
         upstream_tasks = {t.task_id: t for t in task.upstream_list}
@@ -172,6 +168,14 @@ 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 _is_relevant_upstream(upstream: TaskInstance) -> bool:
             """Whether a task instance is a "relevant upstream" of the current task."""
             # Not actually an upstream task.
@@ -265,7 +269,7 @@ 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 = len(setup_upstream_task_instances)  # count of setup tasks upstream of this task
+            upstream_setup = _get_tis_count_from_tasks_list(ti.dag_run, list(setup_upstream_tasks_ids))
         if not any(needs_expansion(t) for t in upstream_tasks.values()):
             upstream = len(upstream_tasks)
             if ti.task.is_teardown:

From 831933f7a03cfcccd1da08f32063952a297e7f81 Mon Sep 17 00:00:00 2001
From: Hussein Awala <hussein@awala.fr>
Date: Wed, 30 Aug 2023 01:01:54 +0200
Subject: [PATCH 22/22] Check if the new method fixes the tests

---
 airflow/ti_deps/deps/trigger_rule_dep.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py
index 39b71ef7d16b0..d73b6dd176694 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -25,7 +25,7 @@
 
 from sqlalchemy import and_, func, or_, select
 
-from airflow.models import DagRun
+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
@@ -36,6 +36,8 @@
     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
 
 
@@ -176,6 +178,18 @@ def _get_tis_count_from_tasks_list(dag_run: DagRun, tasks_ids_list: list[str]) -
                 .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.
@@ -269,7 +283,7 @@ 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_tis_count_from_tasks_list(ti.dag_run, list(setup_upstream_tasks_ids))
+            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)
             if ti.task.is_teardown: