Add reattach_on_existing option to TriggerDagRunOperator

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

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"),
         [

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
 
 

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,18 +1887,19 @@ 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,
                 end_date=datetime.now(tz=timezone.utc),
                 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.