From 5514971b652802a194da1738a0c3bea9a6b1da23 Mon Sep 17 00:00:00 2001 From: luc-pimentel <154371376+luc-pimentel@users.noreply.github.com> Date: Sun, 18 Oct 2026 21:14:57 -0300 Subject: [PATCH 1/2] Test that a task asking to wait for its past depends is not skipped at once A task with depends_on_past=True and wait_for_past_depends_before_skipping=True is documented to stay unscheduled until its previous run's task instance has succeeded or been skipped, and only then be skipped. Today the task attribute is never read: SkipMixin lets the execution API skip the task at once, and the scheduler deps only honour the DepContext flag, which nothing sets. These tests fail on main and pass with the fix: the scheduling decision for a short-circuit and for a branch decision, the SkipMixin lists, the three deps reading the task attribute, and the serialized mapped operator exposing it. --- airflow-core/tests/unit/models/test_dagrun.py | 81 ++++++++++++++++++- .../serialization/test_dag_serialization.py | 12 +++ .../deps/test_not_previously_skipped_dep.py | 38 +++++++++ .../unit/ti_deps/deps/test_prev_dagrun_dep.py | 14 ++++ .../ti_deps/deps/test_trigger_rule_dep.py | 66 +++++++++++++++ .../tests/task_sdk/bases/test_skipmixin.py | 81 +++++++++++++++++++ 6 files changed, 290 insertions(+), 2 deletions(-) diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 5c191505666ca..3ab0958366809 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/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"), [ diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index d40f36ebf6d33..9f3d983e58856 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/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): diff --git a/airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py index 0c0fa53b5245f..9d1a97791af1a 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py +++ b/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 diff --git a/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py index 59a2274c5feb8..107a4e565f8e5 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py @@ -333,3 +333,17 @@ def test_dagrun_dep(mock_get_previous_dagrun, mock_get_previous_scheduled_dagrun ti.xcom_push.assert_called_with(key="past_depends_met", value=True) 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) diff --git a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py index 27fb20d1b785b..5369e5b904740 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py +++ b/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( diff --git a/task-sdk/tests/task_sdk/bases/test_skipmixin.py b/task-sdk/tests/task_sdk/bases/test_skipmixin.py index 9518270e2294a..ebdb747159663 100644 --- a/task-sdk/tests/task_sdk/bases/test_skipmixin.py +++ b/task-sdk/tests/task_sdk/bases/test_skipmixin.py @@ -96,6 +96,87 @@ def test_skip_without_task_id_does_not_push_xcom(self): ti.xcom_push.assert_not_called() + def test_skip_leaves_task_waiting_for_past_depends_to_the_scheduler(self): + """skip() leaves a direct downstream task that waits for its past depends out of the state change.""" + mixin = SkipMixin() + mixin.task_id = "gate" + mixin.downstream_task_ids = {"skip_now", "flag_only", "wait"} + + skip_now = MagicMock( + spec=BaseOperator, + task_id="skip_now", + depends_on_past=True, + wait_for_past_depends_before_skipping=False, + ) + flag_only = MagicMock( + spec=BaseOperator, + task_id="flag_only", + depends_on_past=False, + wait_for_past_depends_before_skipping=True, + ) + wait = MagicMock( + spec=BaseOperator, + task_id="wait", + depends_on_past=True, + wait_for_past_depends_before_skipping=True, + ) + # Reached through "wait", not a direct downstream task: skipped at once as before. + indirect = MagicMock( + spec=BaseOperator, + task_id="indirect", + depends_on_past=True, + wait_for_past_depends_before_skipping=True, + ) + ti = Mock(spec=RuntimeTaskInstanceProtocol, map_index=-1) + + with pytest.raises(DownstreamTasksSkipped) as exc_info: + mixin.skip(ti=ti, tasks=[skip_now, flag_only, wait, indirect]) + + # The decision still lists "wait", so NotPreviouslySkippedDep skips it once its past depends are met. + ti.xcom_push.assert_called_once_with( + key=XCOM_SKIPMIXIN_KEY, + value={XCOM_SKIPMIXIN_SKIPPED: ["skip_now", "flag_only", "wait", "indirect"]}, + ) + assert exc_info.value.tasks == ["skip_now", "flag_only", "indirect"] + + def test_skip_all_except_leaves_task_waiting_for_past_depends_to_the_scheduler(self): + """skip_all_except() leaves a task that waits for its past depends out of the state change.""" + mixin = SkipMixin() + mixin.downstream_task_ids = {"down1", "down2", "down3"} + + downstream1 = MagicMock( + spec=BaseOperator, + task_id="down1", + depends_on_past=True, + wait_for_past_depends_before_skipping=True, + ) + downstream2 = MagicMock( + spec=BaseOperator, + task_id="down2", + depends_on_past=True, + wait_for_past_depends_before_skipping=False, + ) + downstream3 = MagicMock(spec=BaseOperator, task_id="down3") + downstream3.get_flat_relative_ids.return_value = set() + + mock_task = MagicMock(spec=BaseOperator) + mock_task.downstream_list = [downstream1, downstream2, downstream3] + mock_dag = MagicMock(spec=DAG) + mock_dag.task_ids = ["gate", "down1", "down2", "down3"] + mock_dag.get_task.return_value = downstream3 + mock_task.dag = mock_dag + + ti = Mock(spec=RuntimeTaskInstanceProtocol, map_index=-1, task=mock_task) + + with pytest.raises(DownstreamTasksSkipped) as exc_info: + mixin.skip_all_except(ti=ti, branch_task_ids="down3") + + assert exc_info.value.tasks == [("down2", -1)] + ti.xcom_push.assert_called_once_with( + key=XCOM_SKIPMIXIN_KEY, + value={XCOM_SKIPMIXIN_FOLLOWED: ["down3"]}, + ) + def test_skip_all_except_with_none_skips_all(self): """skip_all_except(None) should skip all downstream tasks.""" mixin = SkipMixin() From 99b9a77406213c4ab5867847be310e8ee598d66a Mon Sep 17 00:00:00 2001 From: luc-pimentel <154371376+luc-pimentel@users.noreply.github.com> Date: Sun, 18 Oct 2026 21:44:48 -0300 Subject: [PATCH 2/2] Honour wait_for_past_depends_before_skipping when a task is told to skip A task with depends_on_past=True and wait_for_past_depends_before_skipping=True was skipped at once when an upstream ShortCircuitOperator or branch operator decided to skip it, even though its previous run's task instance had failed. Two things stood between the parameter and the scheduler. SkipMixin raised DownstreamTasksSkipped for every direct downstream task, so the execution API wrote SKIPPED before any dependency ran. And the scheduler deps only honoured DepContext.wait_for_past_depends_before_skipping, which nothing sets in Airflow 3; the task attribute was never read, and SerializedMappedOperator did not expose it. SkipMixin now leaves the direct downstream tasks that set both flags out of the state change, keeping them in the XCom decision, so NotPreviouslySkippedDep decides. The deps read the task attribute as well as the context flag, and PrevDagrunDep writes the past-depends-met mark through the scheduler's session instead of a fresh one that commits inside the scheduling loop. A task that sets only one of the two flags keeps today's immediate skip, as its docstring says, and tasks reached through another skipped task are still skipped at once. --- .../definitions/mappedoperator.py | 4 ++++ .../deps/not_previously_skipped_dep.py | 4 +++- .../airflow/ti_deps/deps/prev_dagrun_dep.py | 23 ++++++++++-------- .../airflow/ti_deps/deps/trigger_rule_dep.py | 12 +++++----- .../unit/ti_deps/deps/test_prev_dagrun_dep.py | 4 ++-- task-sdk/src/airflow/sdk/bases/skipmixin.py | 24 +++++++++++++++++-- 6 files changed, 50 insertions(+), 21 deletions(-) diff --git a/airflow-core/src/airflow/serialization/definitions/mappedoperator.py b/airflow-core/src/airflow/serialization/definitions/mappedoperator.py index 288f1673136b8..8ec4065c6787c 100644 --- a/airflow-core/src/airflow/serialization/definitions/mappedoperator.py +++ b/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") diff --git a/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py b/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py index 80c3380f87745..41abe26d97200 100644 --- a/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py +++ b/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 diff --git a/airflow-core/src/airflow/ti_deps/deps/prev_dagrun_dep.py b/airflow-core/src/airflow/ti_deps/deps/prev_dagrun_dep.py index 1762abb49cd1f..6814fd6616d38 100644 --- a/airflow-core/src/airflow/ti_deps/deps/prev_dagrun_dep.py +++ b/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) diff --git a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py index 9ccc4adcb133b..bcd5c3e2cb9b5 100644 --- a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py +++ b/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( diff --git a/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py index 107a4e565f8e5..1c7a6f601eaae 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py @@ -330,7 +330,7 @@ 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() @@ -346,4 +346,4 @@ def test_past_depends_met_xcom_sent_when_the_task_asks_to_wait(): 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) + ti.xcom_push.assert_called_once_with(key="past_depends_met", value=True, session=ANY) diff --git a/task-sdk/src/airflow/sdk/bases/skipmixin.py b/task-sdk/src/airflow/sdk/bases/skipmixin.py index f4dd25f5f84ab..f759635136059 100644 --- a/task-sdk/src/airflow/sdk/bases/skipmixin.py +++ b/task-sdk/src/airflow/sdk/bases/skipmixin.py @@ -55,6 +55,24 @@ def _set_state_to_skipped( if tasks and map_index == -1: raise DownstreamTasksSkipped(tasks=tasks) + def _tasks_to_skip_now(self, tasks: Sequence[Operator]) -> list[Operator]: + """ + Drop the direct downstream tasks that wait for their past depends before being skipped. + + They stay in the XCom decision, so the scheduler skips them once the previous run's task + instance has succeeded or been skipped (``NotPreviouslySkippedDep``). + """ + direct_downstream_ids: set[str] = getattr(self, "downstream_task_ids", set()) + return [ + t + for t in tasks + if not ( + t.task_id in direct_downstream_ids + and t.depends_on_past + and t.wait_for_past_depends_before_skipping + ) + ] + def skip( self, ti: RuntimeTaskInstanceProtocol, @@ -79,7 +97,7 @@ def skip( value={XCOM_SKIPMIXIN_SKIPPED: task_ids_list}, ) - self._set_state_to_skipped(task_ids_list, ti.map_index) + self._set_state_to_skipped([t.task_id for t in self._tasks_to_skip_now(task_list)], ti.map_index) def skip_all_except( self, @@ -156,7 +174,9 @@ def skip_all_except( branch_task_id_set.update(dag.get_task(branch_task_id).get_flat_relative_ids(upstream=False)) skip_tasks = [ - (t.task_id, ti.map_index) for t in downstream_tasks if t.task_id not in branch_task_id_set + (t.task_id, ti.map_index) + for t in self._tasks_to_skip_now(downstream_tasks) + if t.task_id not in branch_task_id_set ] follow_task_ids = [t.task_id for t in downstream_tasks if t.task_id in branch_task_id_set]