diff --git a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py index 60d96cbe3f495..987ef1946588f 100644 --- a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py +++ b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py @@ -139,6 +139,10 @@ class TriggerDagRunOperator(BaseOperator): DAG run conf is immutable and will not be reset on rerun of an existing DAG run. When reset_dag_run=False and dag run exists, DagRunAlreadyExists will be raised. When reset_dag_run=True and dag run exists, existing DAG run will be cleared to rerun. + :param reattach_on_existing: When ``True`` and the Dag run already exists, attach to the existing + Dag run instead of raising ``DagRunAlreadyExists``. If ``wait_for_completion=True``, the operator + waits for the existing Dag run to reach an allowed or failed state. ``reset_dag_run`` and + ``skip_when_already_exists`` take precedence over this option. :param wait_for_completion: Whether or not wait for DAG run completion. (default: False) :param poke_interval: Poke interval to check DAG run status when wait_for_completion=True. (default: 60) @@ -169,6 +173,7 @@ class TriggerDagRunOperator(BaseOperator): "conf", "wait_for_completion", "skip_when_already_exists", + "reattach_on_existing", ) attributes_not_supported_in_airflow_2 = { @@ -190,6 +195,7 @@ def __init__( logical_date: str | datetime.datetime | None | ArgNotSet = NOTSET, run_after: str | datetime.datetime | None | ArgNotSet = NOTSET, reset_dag_run: bool = False, + reattach_on_existing: bool = False, wait_for_completion: bool = False, poke_interval: int = 60, allowed_states: list[str | DagRunState] | None = None, @@ -206,6 +212,7 @@ def __init__( self.trigger_run_id = trigger_run_id self.conf = conf self.reset_dag_run = reset_dag_run + self.reattach_on_existing = reattach_on_existing self.wait_for_completion = wait_for_completion self.poke_interval = poke_interval if allowed_states: @@ -319,6 +326,12 @@ def _trigger_dag_af_3(self, context, run_id, parsed_logical_date, parsed_run_aft ) parameters = inspect.signature(DagRunTriggerException.__init__).parameters + if "reattach_on_existing" in parameters: + kwargs_accepted["reattach_on_existing"] = self.reattach_on_existing + elif self.reattach_on_existing: + raise NotImplementedError( + "Setting `reattach_on_existing` requires Airflow 3.x task execution support." + ) if self.note and "note" in parameters: kwargs_accepted["note"] = self.note @@ -383,7 +396,13 @@ def _trigger_dag_af_2(self, context, run_id, parsed_logical_date): raise AirflowSkipException( "Skipping due to skip_when_already_exists is set to True and DagRunAlreadyExists" ) - raise e + if self.reattach_on_existing: + dag_run = e.dag_run + self.log.info( + "Reattaching to existing Dag run %s on %s", dag_run.run_id, self.trigger_dag_id + ) + else: + raise e if dag_run is None: raise RuntimeError("The dag_run should be set here!") # Store the run id from the dag run (either created or found above) to diff --git a/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py b/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py index 2bce35574dc01..3cc70437a71d4 100644 --- a/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py +++ b/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py @@ -30,8 +30,8 @@ from airflow.models.dagrun import DagRun from airflow.models.log import Log from airflow.models.taskinstance import TaskInstance -from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred, conf -from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator +from airflow.providers.common.compat.sdk import AirflowException, AirflowSkipException, TaskDeferred, conf +from airflow.providers.standard.operators.trigger_dagrun import XCOM_RUN_ID, TriggerDagRunOperator from airflow.providers.standard.triggers.external_task import DagStateTrigger from airflow.utils.session import create_session from airflow.utils.state import DagRunState, TaskInstanceState @@ -133,6 +133,7 @@ def test_trigger_dagrun_with_run_after(self): assert exc_info.value.logical_date is None assert exc_info.value.reset_dag_run is False assert exc_info.value.skip_when_already_exists is False + assert exc_info.value.reattach_on_existing is False assert exc_info.value.wait_for_completion is False assert exc_info.value.allowed_states == [DagRunState.SUCCESS] assert exc_info.value.failed_states == [DagRunState.FAILED] @@ -168,6 +169,7 @@ def test_trigger_dagrun(self): assert exc_info.value.logical_date is not None assert exc_info.value.reset_dag_run is False assert exc_info.value.skip_when_already_exists is False + assert exc_info.value.reattach_on_existing is False assert exc_info.value.wait_for_completion is False assert exc_info.value.allowed_states == [DagRunState.SUCCESS] assert exc_info.value.failed_states == [DagRunState.FAILED] @@ -245,6 +247,33 @@ def test_trigger_dagrun_custom_run_id(self): assert exc_info.value.dag_run_id == "custom_run_id" + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") + @pytest.mark.parametrize( + ("reset_dag_run", "skip_when_already_exists"), + [ + pytest.param(False, False, id="reattach-only"), + pytest.param(True, False, id="reset-precedence"), + pytest.param(False, True, id="skip-precedence"), + ], + ) + def test_trigger_dagrun_reattach_on_existing(self, reset_dag_run, skip_when_already_exists): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="custom_run_id", + reset_dag_run=reset_dag_run, + skip_when_already_exists=skip_when_already_exists, + reattach_on_existing=True, + ) + + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) + + assert exc_info.value.dag_run_id == "custom_run_id" + assert exc_info.value.reset_dag_run is reset_dag_run + assert exc_info.value.skip_when_already_exists is skip_when_already_exists + assert exc_info.value.reattach_on_existing is True + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") def test_trigger_dagrun_with_logical_date(self): """Test TriggerDagRunOperator with custom logical_date.""" @@ -902,6 +931,152 @@ def test_trigger_dagrun_with_skip_when_already_exists(self, dag_maker): task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) assert dr.get_task_instance("test_task").state == TaskInstanceState.SKIPPED + @mock.patch(f"{TRIGGER_OP_PATH}.SerializedDagModel.get_dag", autospec=True) + @mock.patch(f"{TRIGGER_OP_PATH}.time.sleep", autospec=True) + @mock.patch(f"{TRIGGER_OP_PATH}.trigger_dag", autospec=True) + def test_trigger_dagrun_reattach_on_existing_waits_for_allowed_state( + self, mock_trigger_dag, mock_sleep, mock_get_dag, dag_maker + ): + with dag_maker(TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="dummy_run_id", + logical_date=DEFAULT_DATE, + reset_dag_run=False, + reattach_on_existing=True, + wait_for_completion=True, + poke_interval=10, + ) + + existing_dag_run = mock.MagicMock(spec=DagRun) + existing_dag_run.dag_id = TRIGGERED_DAG_ID + existing_dag_run.run_id = "dummy_run_id" + existing_dag_run.logical_date = DEFAULT_DATE + existing_dag_run.state = DagRunState.SUCCESS + mock_trigger_dag.side_effect = DagRunAlreadyExists(existing_dag_run) + ti = mock.MagicMock(spec=TaskInstance) + + task._trigger_dag_af_2({"task_instance": ti}, "dummy_run_id", DEFAULT_DATE) + + ti.xcom_push.assert_called_once_with(key=XCOM_RUN_ID, value="dummy_run_id") + existing_dag_run.refresh_from_db.assert_called_once_with() + mock_get_dag.assert_not_called() + mock_sleep.assert_called_once_with(10) + + @mock.patch(f"{TRIGGER_OP_PATH}.time.sleep", autospec=True) + @mock.patch(f"{TRIGGER_OP_PATH}.trigger_dag", autospec=True) + def test_trigger_dagrun_reattach_on_existing_failed_state(self, mock_trigger_dag, mock_sleep, dag_maker): + with dag_maker(TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="dummy_run_id", + logical_date=DEFAULT_DATE, + reset_dag_run=False, + reattach_on_existing=True, + wait_for_completion=True, + poke_interval=10, + ) + + existing_dag_run = mock.MagicMock(spec=DagRun) + existing_dag_run.dag_id = TRIGGERED_DAG_ID + existing_dag_run.run_id = "dummy_run_id" + existing_dag_run.logical_date = DEFAULT_DATE + existing_dag_run.state = DagRunState.FAILED + mock_trigger_dag.side_effect = DagRunAlreadyExists(existing_dag_run) + + with pytest.raises(AirflowException, match=f"{TRIGGERED_DAG_ID} failed"): + task._trigger_dag_af_2( + {"task_instance": mock.MagicMock(spec=TaskInstance)}, "dummy_run_id", DEFAULT_DATE + ) + + existing_dag_run.refresh_from_db.assert_called_once_with() + mock_sleep.assert_called_once_with(10) + + @mock.patch(f"{TRIGGER_OP_PATH}.trigger_dag", autospec=True) + def test_trigger_dagrun_reattach_on_existing_defers(self, mock_trigger_dag, dag_maker): + with dag_maker(TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="dummy_run_id", + logical_date=DEFAULT_DATE, + reset_dag_run=False, + reattach_on_existing=True, + wait_for_completion=True, + deferrable=True, + poke_interval=5, + ) + + existing_dag_run = mock.MagicMock(spec=DagRun) + existing_dag_run.dag_id = TRIGGERED_DAG_ID + existing_dag_run.run_id = "dummy_run_id" + existing_dag_run.logical_date = DEFAULT_DATE + mock_trigger_dag.side_effect = DagRunAlreadyExists(existing_dag_run) + mock_task_defer = mock.MagicMock(side_effect=task.defer) + + with mock.patch.object(TriggerDagRunOperator, "defer", mock_task_defer), pytest.raises(TaskDeferred): + task._trigger_dag_af_2( + {"task_instance": mock.MagicMock(spec=TaskInstance)}, "dummy_run_id", DEFAULT_DATE + ) + + trigger = mock_task_defer.call_args.kwargs["trigger"] + assert trigger.run_ids == ["dummy_run_id"] + assert trigger.execution_dates == [DEFAULT_DATE] + + @mock.patch(f"{TRIGGER_OP_PATH}.trigger_dag", autospec=True) + def test_trigger_dagrun_skip_precedes_reattach_on_existing(self, mock_trigger_dag, dag_maker): + with dag_maker(TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="dummy_run_id", + reset_dag_run=False, + skip_when_already_exists=True, + reattach_on_existing=True, + ) + + existing_dag_run = mock.MagicMock(spec=DagRun) + existing_dag_run.dag_id = TRIGGERED_DAG_ID + existing_dag_run.run_id = "dummy_run_id" + mock_trigger_dag.side_effect = DagRunAlreadyExists(existing_dag_run) + + with pytest.raises(AirflowSkipException): + task._trigger_dag_af_2( + {"task_instance": mock.MagicMock(spec=TaskInstance)}, "dummy_run_id", DEFAULT_DATE + ) + + @mock.patch(f"{TRIGGER_OP_PATH}.DagModel.get_current", autospec=True) + @mock.patch(f"{TRIGGER_OP_PATH}.SerializedDagModel.get_dag", autospec=True) + @mock.patch(f"{TRIGGER_OP_PATH}.trigger_dag", autospec=True) + def test_trigger_dagrun_reset_precedes_reattach_on_existing( + self, mock_trigger_dag, mock_get_dag, mock_get_current, dag_maker + ): + with dag_maker(TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="dummy_run_id", + logical_date=DEFAULT_DATE, + reset_dag_run=True, + reattach_on_existing=True, + ) + + existing_dag_run = mock.MagicMock(spec=DagRun) + existing_dag_run.dag_id = TRIGGERED_DAG_ID + existing_dag_run.run_id = "dummy_run_id" + existing_dag_run.logical_date = DEFAULT_DATE + mock_trigger_dag.side_effect = DagRunAlreadyExists(existing_dag_run) + mock_get_current.return_value = mock.MagicMock(spec=DagModel) + dag = mock_get_dag.return_value + ti = mock.MagicMock(spec=TaskInstance) + + task._trigger_dag_af_2({"task_instance": ti}, "dummy_run_id", DEFAULT_DATE) + + dag.clear.assert_called_once_with(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + ti.xcom_push.assert_called_once_with(key=XCOM_RUN_ID, value="dummy_run_id") + @pytest.mark.parametrize( ("trigger_run_id", "trigger_logical_date", "expected_dagruns_count"), [ diff --git a/task-sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py index 6f43d5421ecf2..805b46c7bd352 100644 --- a/task-sdk/src/airflow/sdk/exceptions.py +++ b/task-sdk/src/airflow/sdk/exceptions.py @@ -309,6 +309,7 @@ def __init__( failed_states: list[str], poke_interval: int, deferrable: bool, + reattach_on_existing: bool = False, note: str | None = None, ): super().__init__() @@ -324,6 +325,7 @@ def __init__( self.failed_states = failed_states self.poke_interval = poke_interval self.deferrable = deferrable + self.reattach_on_existing = reattach_on_existing self.note = note diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 94363b984b7f6..66a319e8dea4b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1872,7 +1872,10 @@ def _handle_trigger_dag_run( ), ) - if isinstance(comms_msg, ErrorResponse) and comms_msg.error == ErrorType.DAGRUN_ALREADY_EXISTS: + dag_run_already_exists = ( + isinstance(comms_msg, ErrorResponse) and comms_msg.error == ErrorType.DAGRUN_ALREADY_EXISTS + ) + if dag_run_already_exists: if drte.skip_when_already_exists: log.info( "Dag Run already exists, skipping task as skip_when_already_exists is set to True.", @@ -1884,7 +1887,8 @@ def _handle_trigger_dag_run( rendered_map_index=ti.rendered_map_index, ) state = TaskInstanceState.SKIPPED - else: + return msg, state + if not drte.reattach_on_existing: log.error("Dag Run already exists, marking task as failed.", dag_id=drte.trigger_dag_id) msg = TaskState( state=TaskInstanceState.FAILED, @@ -1892,10 +1896,10 @@ def _handle_trigger_dag_run( rendered_map_index=ti.rendered_map_index, ) state = TaskInstanceState.FAILED - - return msg, state - - log.info("Dag Run triggered successfully.", trigger_dag_id=drte.trigger_dag_id) + return msg, state + log.info("Dag Run already exists, reattaching to it.", dag_id=drte.trigger_dag_id) + else: + log.info("Dag Run triggered successfully.", trigger_dag_id=drte.trigger_dag_id) # Store the run id from the dag run (either created or found above) to # be used when creating the extra link on the webserver. diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 9d15ac40c6967..3bbbeab3aba54 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -5048,6 +5048,171 @@ def test_handle_trigger_dag_run_conflict( ] mock_supervisor_comms.assert_has_calls(expected_calls) + @time_machine.travel("2025-01-01 00:00:00", tick=False) + def test_handle_trigger_dag_run_conflict_reattach(self, create_runtime_ti, mock_supervisor_comms): + from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator + + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id="test_dag", + trigger_run_id="test_run_id", + reattach_on_existing=True, + ) + ti = create_runtime_ti( + dag_id="test_handle_trigger_dag_run_conflict_reattach", run_id="test_run", task=task + ) + + def _send_side_effect(*args, **kwargs): + msg = kwargs.get("msg") + if msg is None and args: + msg = args[0] + if isinstance(msg, TriggerDagRun): + return ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + return None + + mock_supervisor_comms.send.side_effect = _send_side_effect + + state, msg, _ = run(ti, ti.get_template_context(), mock.MagicMock()) + + assert state == TaskInstanceState.SUCCESS + assert msg.state == TaskInstanceState.SUCCESS + sent_messages = [ + call.args[0] if call.args else call.kwargs["msg"] + for call in mock_supervisor_comms.send.call_args_list + ] + trigger_msg = next(msg for msg in sent_messages if isinstance(msg, TriggerDagRun)) + assert trigger_msg.dag_id == "test_dag" + assert trigger_msg.run_id == "test_run_id" + assert trigger_msg.reset_dag_run is False + xcom_msg = next( + msg for msg in sent_messages if isinstance(msg, SetXCom) and msg.key == "trigger_run_id" + ) + assert xcom_msg.value == "test_run_id" + assert xcom_msg.dag_id == "test_handle_trigger_dag_run_conflict_reattach" + assert xcom_msg.task_id == "test_task" + assert xcom_msg.run_id == "test_run" + + @time_machine.travel("2025-01-01 00:00:00", tick=False) + def test_handle_trigger_dag_run_conflict_skip_precedes_reattach( + self, create_runtime_ti, mock_supervisor_comms + ): + from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator + + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id="test_dag", + trigger_run_id="test_run_id", + skip_when_already_exists=True, + reattach_on_existing=True, + ) + ti = create_runtime_ti( + dag_id="test_handle_trigger_dag_run_conflict_skip_precedes_reattach", run_id="test_run", task=task + ) + + mock_supervisor_comms.send.return_value = ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + state, msg, _ = run(ti, ti.get_template_context(), mock.MagicMock()) + + assert state == TaskInstanceState.SKIPPED + assert msg.state == TaskInstanceState.SKIPPED + + @pytest.mark.parametrize( + ("target_dr_state", "expected_task_state"), + [ + pytest.param(DagRunState.SUCCESS, TaskInstanceState.SUCCESS, id="success"), + pytest.param(DagRunState.FAILED, TaskInstanceState.FAILED, id="failed"), + ], + ) + @time_machine.travel("2025-01-01 00:00:00", tick=False) + def test_handle_trigger_dag_run_conflict_reattach_wait_for_completion( + self, target_dr_state, expected_task_state, create_runtime_ti, mock_supervisor_comms + ): + from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator + + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id="test_dag", + trigger_run_id="test_run_id", + poke_interval=5, + wait_for_completion=True, + reattach_on_existing=True, + ) + ti = create_runtime_ti( + dag_id="test_handle_trigger_dag_run_conflict_reattach_wait", run_id="test_run", task=task + ) + dag_run_states = iter([DagRunState.RUNNING, target_dr_state]) + + def _send_side_effect(*args, **kwargs): + msg = kwargs.get("msg") + if msg is None and args: + msg = args[0] + if isinstance(msg, TriggerDagRun): + return ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + if isinstance(msg, GetDagRunState): + return DagRunStateResult(state=next(dag_run_states)) + return None + + mock_supervisor_comms.send.side_effect = _send_side_effect + + with mock.patch("time.sleep", autospec=True, return_value=None): + state, msg, _ = run(ti, ti.get_template_context(), mock.MagicMock()) + + assert state == expected_task_state + assert msg.state == expected_task_state + sent_messages = [ + call.args[0] if call.args else call.kwargs["msg"] + for call in mock_supervisor_comms.send.call_args_list + ] + xcom_msg = next( + msg for msg in sent_messages if isinstance(msg, SetXCom) and msg.key == "trigger_run_id" + ) + assert xcom_msg.value == "test_run_id" + assert xcom_msg.dag_id == "test_handle_trigger_dag_run_conflict_reattach_wait" + assert xcom_msg.task_id == "test_task" + assert xcom_msg.run_id == "test_run" + dag_run_state_messages = [msg for msg in sent_messages if isinstance(msg, GetDagRunState)] + assert dag_run_state_messages == [ + GetDagRunState(dag_id="test_dag", run_id="test_run_id"), + GetDagRunState(dag_id="test_dag", run_id="test_run_id"), + ] + + @time_machine.travel("2025-01-01 00:00:00", tick=False) + def test_handle_trigger_dag_run_conflict_reattach_deferred( + self, create_runtime_ti, mock_supervisor_comms + ): + from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator + + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id="test_dag", + trigger_run_id="test_run_id", + wait_for_completion=True, + deferrable=True, + reattach_on_existing=True, + poke_interval=5, + ) + ti = create_runtime_ti( + dag_id="test_handle_trigger_dag_run_conflict_reattach_deferred", run_id="test_run", task=task + ) + + def _send_side_effect(*args, **kwargs): + msg = kwargs.get("msg") + if msg is None and args: + msg = args[0] + if isinstance(msg, TriggerDagRun): + return ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + return None + + mock_supervisor_comms.send.side_effect = _send_side_effect + + state, msg, _ = run(ti, ti.get_template_context(), mock.MagicMock()) + + assert state == TaskInstanceState.DEFERRED + assert isinstance(msg, DeferTask) + assert msg.trigger_kwargs["dag_id"] == "test_dag" + assert msg.trigger_kwargs["run_ids"] == ["test_run_id"] + assert msg.trigger_kwargs["execution_dates"] is None + assert msg.trigger_kwargs["poll_interval"] == 5 + @time_machine.travel("2025-01-01 00:00:00", tick=False) def test_handle_trigger_dag_run_reraises_original_error(self, create_runtime_ti, mock_supervisor_comms): """