Honor wait_for_past_depends_before_skipping when an upstream skips

airflow-core/src/airflow/serialization/definitions/mappedoperator.py

@@ -223,6 +223,10 @@ def ignore_first_depends_on_past(self) -> bool:
     def wait_for_downstream(self) -> bool:
         return self._get_partial_kwargs_or_operator_default("wait_for_downstream")
 
+    @property
+    def wait_for_past_depends_before_skipping(self) -> bool:
+        return self._get_partial_kwargs_or_operator_default("wait_for_past_depends_before_skipping")
+
     @property
     def retries(self) -> int:
         return self._get_partial_kwargs_or_operator_default("retries")

airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py

@@ -140,7 +140,9 @@ def _skip(
         # If the parent SkipMixin has run, and the XCom result stored indicates this
         # ti should be skipped, set ti.state to SKIPPED and fail the rule so that the
         # ti does not execute.
-        if dep_context.wait_for_past_depends_before_skipping:
+        if TYPE_CHECKING:
+            assert ti.task
+        if dep_context.wait_for_past_depends_before_skipping or ti.task.wait_for_past_depends_before_skipping:
             if not is_past_depends_met(ti, session=session):
                 yield self._failing_status(reason="Task should be skipped but the past depends are not met")
                 return

airflow-core/src/airflow/ti_deps/deps/prev_dagrun_dep.py

@@ -64,9 +64,12 @@ class PrevDagrunDep(BaseTIDep):
     IS_TASK_DEP = True
 
     @staticmethod
-    def _push_past_deps_met_xcom_if_needed(ti: TI, dep_context):
-        if dep_context.wait_for_past_depends_before_skipping:
-            ti.xcom_push(key=PAST_DEPENDS_MET, value=True)
+    def _push_past_deps_met_xcom_if_needed(ti: TI, dep_context, *, session: Session):
+        if TYPE_CHECKING:
+            assert ti.task
+        if dep_context.wait_for_past_depends_before_skipping or ti.task.wait_for_past_depends_before_skipping:
+            # The scheduler's session: a fresh one would commit inside the scheduling loop.
+            ti.xcom_push(key=PAST_DEPENDS_MET, value=True, session=session)
 
     @staticmethod
     def _has_tis(dagrun: DagRun, task_id: str, *, session: Session) -> bool:
@@ -148,13 +151,13 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session):
         if TYPE_CHECKING:
             assert ti.task
         if dep_context.ignore_depends_on_past:
-            self._push_past_deps_met_xcom_if_needed(ti, dep_context)
+            self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session)
             reason = "The context specified that the state of past DAGs could be ignored."
             yield self._passing_status(reason=reason)
             return
 
         if not ti.task.depends_on_past:
-            self._push_past_deps_met_xcom_if_needed(ti, dep_context)
+            self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session)
             yield self._passing_status(reason="The task did not have depends_on_past set.")
             return
 
@@ -171,7 +174,7 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session):
                 return
 
         if not dr:
-            self._push_past_deps_met_xcom_if_needed(ti, dep_context)
+            self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session)
             yield self._passing_status(reason="This task instance does not belong to a DAG.")
             return
 
@@ -183,7 +186,7 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session):
             last_dagrun = DagRun.get_previous_dagrun(dr, session=session)
         # First ever run for this DAG.
         if not last_dagrun:
-            self._push_past_deps_met_xcom_if_needed(ti, dep_context)
+            self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session)
             yield self._passing_status(reason="This task instance was the first task instance for its task.")
             return
 
@@ -194,14 +197,14 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session):
             and last_dagrun.logical_date is not None
             and last_dagrun.logical_date < ti.task.start_date
         ):
-            self._push_past_deps_met_xcom_if_needed(ti, dep_context)
+            self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session)
             yield self._passing_status(reason="This task instance was the first task instance for its task.")
             return
 
         if not self._has_tis(last_dagrun, ti.task_id, session=session):
             if ti.task.ignore_first_depends_on_past:
                 if not self._has_any_prior_tis(ti, session=session):
-                    self._push_past_deps_met_xcom_if_needed(ti, dep_context)
+                    self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session)
                     yield self._passing_status(
                         reason="ignore_first_depends_on_past is true for this task "
                         "and it is the first task instance for its task."
@@ -234,4 +237,4 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session):
             )
             return
 
-        self._push_past_deps_met_xcom_if_needed(ti, dep_context)
+        self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session)

airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py

@@ -366,9 +366,9 @@ def _evaluate_setup_constraint(
                         new_state = TaskInstanceState.REMOVED
 
             if new_state is not None:
-                if (
-                    new_state == TaskInstanceState.SKIPPED
-                    and dep_context.wait_for_past_depends_before_skipping
+                if new_state == TaskInstanceState.SKIPPED and (
+                    dep_context.wait_for_past_depends_before_skipping
+                    or task.wait_for_past_depends_before_skipping
                 ):
                     if not is_past_depends_met(ti, session=session):
                         yield (
@@ -524,9 +524,9 @@ def _evaluate_direct_relatives() -> Iterator[TIDepStatus]:
                         # 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
+                if new_state == TaskInstanceState.SKIPPED and (
+                    dep_context.wait_for_past_depends_before_skipping
+                    or task.wait_for_past_depends_before_skipping
                 ):
                     if not is_past_depends_met(ti, session=session):
                         yield self._failing_status(

airflow-core/tests/unit/models/test_dagrun.py

@@ -81,10 +81,14 @@
 from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.trigger import Trigger
 from airflow.models.variable import Variable
-from airflow.models.xcom import XComModelV2
+from airflow.models.xcom import XComModel, XComModelV2
 from airflow.providers.standard.operators.bash import BashOperator
 from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.providers.standard.operators.python import PythonOperator, ShortCircuitOperator
+from airflow.providers.standard.operators.python import (
+    BranchPythonOperator,
+    PythonOperator,
+    ShortCircuitOperator,
+)
 from airflow.sdk import (
     DAG,
     BaseOperator,
@@ -103,6 +107,11 @@
 from airflow.settings import get_policy_plugin_manager
 from airflow.task.trigger_rule import TriggerRule
 from airflow.ti_deps.dep_context import DepContext
+from airflow.ti_deps.deps.not_previously_skipped_dep import (
+    XCOM_SKIPMIXIN_FOLLOWED,
+    XCOM_SKIPMIXIN_KEY,
+    XCOM_SKIPMIXIN_SKIPPED,
+)
 from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep
 from airflow.triggers.base import StartTriggerArgs
 from airflow.utils.session import create_session
@@ -1284,6 +1293,74 @@ def test_depends_on_past(self, dag_maker, session, prev_ti_state, is_ti_schedula
         schedulable_tis = [ti.task_id for ti in decision.schedulable_tis]
         assert ("test_dop_task" in schedulable_tis) == is_ti_schedulable
 
+    @pytest.mark.parametrize(
+        ("skipmixin_operator", "decision"),
+        [
+            pytest.param(ShortCircuitOperator, {XCOM_SKIPMIXIN_SKIPPED: ["past"]}, id="short_circuit"),
+            pytest.param(BranchPythonOperator, {XCOM_SKIPMIXIN_FOLLOWED: []}, id="branch"),
+        ],
+    )
+    @pytest.mark.parametrize(
+        ("wait", "state_while_past_failed"),
+        [
+            pytest.param(False, TaskInstanceState.SKIPPED, id="skip_now"),
+            pytest.param(True, None, id="wait_for_past_depends"),
+        ],
+    )
+    def test_skip_decision_waits_for_past_depends(
+        self, dag_maker, session, skipmixin_operator, decision, wait, state_while_past_failed
+    ):
+        """A task told to skip waits for its past depends when it asks to, and is skipped once they are met."""
+        with dag_maker(
+            dag_id="test_skip_decision_waits_for_past_depends",
+            schedule=datetime.timedelta(days=1),
+            session=session,
+            serialized=True,
+        ):
+            gate = skipmixin_operator(task_id="gate", python_callable=lambda: None)
+            past = BaseOperator(
+                task_id="past", depends_on_past=True, wait_for_past_depends_before_skipping=wait
+            )
+            gate >> past
+
+        dag_run_1: DagRun = dag_maker.create_dagrun(
+            logical_date=timezone.datetime(2016, 1, 1, 0, 0, 0), run_type=DagRunType.SCHEDULED
+        )
+        dag_run_2: DagRun = dag_maker.create_dagrun(
+            logical_date=timezone.datetime(2016, 1, 2, 0, 0, 0), run_type=DagRunType.SCHEDULED
+        )
+        # Run 1: the gate ran and the past-dependent task failed.
+        dag_run_1.get_task_instance("gate", session=session).set_state(
+            TaskInstanceState.SUCCESS, session=session
+        )
+        past_ti_1 = dag_run_1.get_task_instance("past", session=session)
+        past_ti_1.set_state(TaskInstanceState.FAILED, session=session)
+        # Run 2: the gate decided that "past" should be skipped.
+        dag_run_2.get_task_instance("gate", session=session).set_state(
+            TaskInstanceState.SUCCESS, session=session
+        )
+        XComModel.set(
+            key=XCOM_SKIPMIXIN_KEY,
+            value=decision,
+            dag_id=dag_run_2.dag_id,
+            task_id="gate",
+            run_id=dag_run_2.run_id,
+            session=session,
+        )
+        session.flush()
+
+        dag_run_2.task_instance_scheduling_decisions(session=session)
+        session.flush()
+        assert dag_run_2.get_task_instance("past", session=session).state == state_while_past_failed
+
+        past_ti_1.set_state(TaskInstanceState.SUCCESS, session=session)
+        session.flush()
+        # The past-depends-met mark and the skip may land in separate scheduling passes.
+        for _ in range(2):
+            dag_run_2.task_instance_scheduling_decisions(session=session)
+            session.flush()
+        assert dag_run_2.get_task_instance("past", session=session).state == TaskInstanceState.SKIPPED
+
     @pytest.mark.parametrize(
         ("prev_ti_state", "is_ti_schedulable"),
         [

airflow-core/tests/unit/serialization/test_dag_serialization.py

@@ -2075,6 +2075,18 @@ def test_serialize_mapped_outlets(self):
         assert round_tripped.inlets == []
         assert round_tripped.outlets == []
 
+    def test_serialize_mapped_wait_for_past_depends_before_skipping(self):
+        with DAG(dag_id="d", schedule=None, start_date=datetime.now()):
+            op = MockOperator.partial(
+                task_id="x", depends_on_past=True, wait_for_past_depends_before_skipping=True
+            ).expand(arg1=[1, 2])
+
+        serialized = OperatorSerialization.serialize_mapped_operator(op)
+        round_tripped = OperatorSerialization.deserialize_operator(serialized)
+        assert isinstance(round_tripped, SerializedMappedOperator)
+        assert round_tripped.depends_on_past is True
+        assert round_tripped.wait_for_past_depends_before_skipping is True
+
     @pytest.mark.db_test
     @pytest.mark.parametrize("mapped", [False, True])
     def test_derived_dag_deps_sensor(self, mapped):

airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py

@@ -152,6 +152,44 @@ def test_parent_skip_branch(session, dag_maker):
     assert tis["op2"].state == State.SKIPPED
 
 
+def test_parent_skip_branch_waits_for_past_depends_when_the_task_asks(session, dag_maker):
+    """
+    A BranchPythonOperator does not follow op2, which waits for its past depends before being skipped.
+    NotPreviouslySkippedDep is not met and leaves op2 unchanged until the past depends are met.
+    """
+    start_date = pendulum.datetime(2020, 1, 1)
+    with dag_maker(
+        "test_parent_skip_branch_waits_for_past_depends_dag",
+        schedule=None,
+        start_date=start_date,
+        session=session,
+    ):
+        op1 = BranchPythonOperator(task_id="op1", python_callable=lambda: "op3")
+        op2 = EmptyOperator(task_id="op2", depends_on_past=True, wait_for_past_depends_before_skipping=True)
+        op3 = EmptyOperator(task_id="op3")
+        op1 >> [op2, op3]
+
+    dagrun = dag_maker.create_dagrun(run_type=DagRunType.MANUAL, state=State.RUNNING)
+    tis = {ti.task_id: ti for ti in dagrun.task_instances}
+    run_task_instance(tis["op1"], op1)
+
+    dep = NotPreviouslySkippedDep()
+    (status,) = dep.get_dep_statuses(tis["op2"], DepContext(), session=session)
+    assert "past depends are not met" in status.reason
+    assert tis["op2"].state == State.NONE
+
+    XComModel.set(
+        key=PAST_DEPENDS_MET,
+        value=True,
+        dag_id=dagrun.dag_id,
+        task_id="op2",
+        run_id=dagrun.run_id,
+        session=session,
+    )
+    assert not dep.is_met(tis["op2"], session=session)
+    assert tis["op2"].state == State.SKIPPED
+
+
 def test_parent_not_executed(session, dag_maker):
     """
     A simple DAG with a BranchPythonOperator that does not follow op2. Parent task is not yet

airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py

@@ -330,6 +330,20 @@ def test_dagrun_dep(mock_get_previous_dagrun, mock_get_previous_scheduled_dagrun
 
     assert actual_dep_met == expected_dep_met
     if past_depends_met_xcom_sent:
-        ti.xcom_push.assert_called_with(key="past_depends_met", value=True)
+        ti.xcom_push.assert_called_with(key="past_depends_met", value=True, session=ANY)
     else:
         ti.xcom_push.assert_not_called()
+
+
+def test_past_depends_met_xcom_sent_when_the_task_asks_to_wait():
+    """The task's own wait_for_past_depends_before_skipping sends the mark like the context flag does."""
+    task = BaseOperator(
+        task_id="test_task",
+        dag=DAG("test_dag", schedule=timedelta(days=1), start_date=datetime(2016, 1, 1)),
+        start_date=datetime(2016, 1, 1),
+        wait_for_past_depends_before_skipping=True,
+    )
+    ti = Mock(task=task, task_id=task.task_id, **{"xcom_push.return_value": None})
+
+    assert PrevDagrunDep().is_met(ti=ti, dep_context=DepContext(ignore_depends_on_past=True))
+    ti.xcom_push.assert_called_once_with(key="past_depends_met", value=True, session=ANY)

airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py

@@ -734,6 +734,40 @@ def test_all_success_tr_skip_wait_for_past_depends_before_skipping_past_depends_
                 expected_reason="requires all upstream tasks to have succeeded, but found 1",
             )
 
+    @pytest.mark.parametrize(
+        ("past_depends_met", "expected_ti_state", "expected_reason"),
+        [
+            pytest.param(None, None, "Task should be skipped but the past depends are not met", id="unmet"),
+            pytest.param(
+                True, SKIPPED, "requires all upstream tasks to have succeeded, but found 1", id="met"
+            ),
+        ],
+    )
+    def test_all_success_tr_skip_waits_when_the_task_asks(
+        self, session, get_task_instance, past_depends_met, expected_ti_state, expected_reason
+    ):
+        """The task's own wait_for_past_depends_before_skipping holds the skip like the context flag does."""
+        ti = get_task_instance(
+            TriggerRule.ALL_SUCCESS,
+            success=1,
+            skipped=1,
+            failed=0,
+            removed=0,
+            upstream_failed=0,
+            done=2,
+            normal_tasks=["FakeTaskID", "OtherFakeTaskID"],
+        )
+        ti.task.wait_for_past_depends_before_skipping = True
+        xcom_mock = Mock(return_value=past_depends_met)
+        with mock.patch("airflow.models.taskinstance.TaskInstance.xcom_pull", xcom_mock):
+            _test_trigger_rule(
+                ti=ti,
+                session=session,
+                flag_upstream_failed=True,
+                expected_ti_state=expected_ti_state,
+                expected_reason=expected_reason,
+            )
+
     @pytest.mark.parametrize("flag_upstream_failed", [True, False])
     def test_none_failed_tr_success(self, session, get_task_instance, flag_upstream_failed):
         """
@@ -2355,6 +2389,38 @@ def test_wait_for_past_depends_before_skipping_reads_stored_value(
     m.assert_not_called()
 
 
+@pytest.mark.parametrize(
+    ("past_depends_met", "expected_ti_state", "expect_failure"),
+    [pytest.param(None, None, True, id="unmet"), pytest.param(True, SKIPPED, False, id="met")],
+)
+def test_setup_constraint_waits_when_the_task_asks(
+    session, get_task_instance, past_depends_met, expected_ti_state, expect_failure
+):
+    """Setup task with a skipped upstream task: the task's own flag holds the skip like the context flag does."""
+    ti = get_task_instance(
+        trigger_rule=TriggerRule.ALL_DONE,
+        success=1,
+        skipped=1,
+        failed=0,
+        removed=0,
+        upstream_failed=0,
+        done=2,
+        setup_tasks=["FakeTaskID", "OtherFakeTaskID"],
+    )
+    ti.task.wait_for_past_depends_before_skipping = True
+    xcom_mock = Mock(return_value=past_depends_met)
+    with mock.patch("airflow.models.taskinstance.TaskInstance.xcom_pull", xcom_mock):
+        _test_trigger_rule(
+            ti=ti,
+            session=session,
+            flag_upstream_failed=True,
+            expected_ti_state=expected_ti_state,
+            expected_reason=(
+                "Task should be skipped but the past depends are not met" if expect_failure else ""
+            ),
+        )
+
+
 @pytest.mark.parametrize(("flag_upstream_failed", "expected_ti_state"), [(True, SKIPPED), (False, None)])
 @pytest.mark.need_serialized_dag
 def test_setup_mapped_task_group_finished_upstream_before_expand(