Skip to content

Commit c46395b

Browse files
hwang-cadenthwang-cadent
committed
Fix dynamic dag_id resolution in TriggerDagRunOperator links
- Add XCOM_DAG_ID constant to store resolved dag_id in XCom - Update TriggerDagRunLink.get_link() to check XCom first for dynamic dag_ids - Store resolved dag_id in XCom during execution for both Airflow 2.x and 3.x - Add comprehensive tests for dynamic dag_id link generation - Maintain backward compatibility with existing static dag_id usage - Fix deserialization of logical_date when it's NOTSET Fixes #46402 diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index db79e79..a9f1c3c770 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -1595,6 +1595,11 @@ class OperatorSerialization(DAGNode, BaseSerialization): elif field_name == "resources": return Resources.from_dict(value) if value is not None else None elif field_name.endswith("_date"): + # Check if value is ARG_NOT_SET before trying to deserialize as datetime + if isinstance(value, dict) and value.get(Encoding.TYPE) == DAT.ARG_NOT_SET: + from airflow.serialization.definitions.notset import NOTSET + + return NOTSET return cls._deserialize_datetime(value) if value is not None else None else: # For all other fields, return as-is (strings, ints, bools, etc.) 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 ae3f978..728a1cf 100644 --- a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py +++ b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py @@ -53,6 +53,7 @@ except ImportError: XCOM_LOGICAL_DATE_ISO = "trigger_logical_date_iso" XCOM_RUN_ID = "trigger_run_id" +XCOM_DAG_ID = "trigger_dag_id" if TYPE_CHECKING: @@ -85,21 +86,26 @@ class TriggerDagRunLink(BaseOperatorLink): if TYPE_CHECKING: assert isinstance(operator, TriggerDagRunOperator) - trigger_dag_id = operator.trigger_dag_id - if not AIRFLOW_V_3_0_PLUS: - from airflow.models.renderedtifields import RenderedTaskInstanceFields - from airflow.models.taskinstancekey import TaskInstanceKey as CoreTaskInstanceKey - - core_ti_key = CoreTaskInstanceKey( - dag_id=ti_key.dag_id, - task_id=ti_key.task_id, - run_id=ti_key.run_id, - try_number=ti_key.try_number, - map_index=ti_key.map_index, - ) + # Try to get the resolved dag_id from XCom first (for dynamic dag_ids) + trigger_dag_id = XCom.get_value(ti_key=ti_key, key=XCOM_DAG_ID) + + # Fallback to operator attribute and rendered fields if not in XCom + if not trigger_dag_id: + trigger_dag_id = operator.trigger_dag_id + if not AIRFLOW_V_3_0_PLUS: + from airflow.models.renderedtifields import RenderedTaskInstanceFields + from airflow.models.taskinstancekey import TaskInstanceKey as CoreTaskInstanceKey + + core_ti_key = CoreTaskInstanceKey( + dag_id=ti_key.dag_id, + task_id=ti_key.task_id, + run_id=ti_key.run_id, + try_number=ti_key.try_number, + map_index=ti_key.map_index, + ) - if template_fields := RenderedTaskInstanceFields.get_templated_fields(core_ti_key): - trigger_dag_id: str = template_fields.get("trigger_dag_id", operator.trigger_dag_id) # type: ignore[no-redef] + if template_fields := RenderedTaskInstanceFields.get_templated_fields(core_ti_key): + trigger_dag_id: str = template_fields.get("trigger_dag_id", operator.trigger_dag_id) # type: ignore[no-redef] # Fetch the correct dag_run_id for the triggerED dag which is # stored in xcom during execution of the triggerING task. @@ -203,7 +209,7 @@ class TriggerDagRunOperator(BaseOperator): self.openlineage_inject_parent_info = openlineage_inject_parent_info self.deferrable = deferrable self.logical_date = logical_date - if logical_date is NOTSET: + if isinstance(logical_date, ArgNotSet) or logical_date is NOTSET: self.logical_date = NOTSET elif logical_date is None or isinstance(logical_date, (str, datetime.datetime)): self.logical_date = logical_date @@ -216,7 +222,7 @@ class TriggerDagRunOperator(BaseOperator): raise NotImplementedError("Setting `fail_when_dag_is_paused` not yet supported for Airflow 3.x") def execute(self, context: Context): - if self.logical_date is NOTSET: + if isinstance(self.logical_date, ArgNotSet) or self.logical_date is NOTSET: # If no logical_date is provided we will set utcnow() parsed_logical_date = timezone.utcnow() elif self.logical_date is None or isinstance(self.logical_date, datetime.datetime): @@ -274,6 +280,14 @@ class TriggerDagRunOperator(BaseOperator): def _trigger_dag_af_3(self, context, run_id, parsed_logical_date): from airflow.providers.common.compat.sdk import DagRunTriggerException + # Store the resolved dag_id to XCom for use in the link generation + # This is important for dynamic dag_ids (from XCom or complex templates) + # In Airflow 3.x, context has both "task_instance" and "ti" keys + if "task_instance" in context: + context["task_instance"].xcom_push(key=XCOM_DAG_ID, value=self.trigger_dag_id) + elif "ti" in context: + context["ti"].xcom_push(key=XCOM_DAG_ID, value=self.trigger_dag_id) + raise DagRunTriggerException( trigger_dag_id=self.trigger_dag_id, dag_run_id=run_id, @@ -319,10 +333,11 @@ class TriggerDagRunOperator(BaseOperator): 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 + # Store the run id and dag_id from the dag run (either created or found above) to # be used when creating the extra link on the webserver. ti = context["task_instance"] ti.xcom_push(key=XCOM_RUN_ID, value=dag_run.run_id) + ti.xcom_push(key=XCOM_DAG_ID, value=self.trigger_dag_id) if self.wait_for_completion: # Kick off the deferral process 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 0f8d171..920f38b 100644 --- a/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py +++ b/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py @@ -140,8 +140,10 @@ class TestDagRunOperator: assert task.trigger_run_id == expected_run_id # run_id is saved as attribute @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") - @mock.patch(f"{TRIGGER_OP_PATH}.XCom.get_one") - def test_extra_operator_link(self, mock_xcom_get_one, dag_maker): + @mock.patch("airflow.providers.standard.operators.trigger_dagrun.XCom.get_value") + def test_extra_operator_link(self, mock_xcom_get_value, dag_maker): + from airflow.providers.standard.operators.trigger_dagrun import XCOM_RUN_ID + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): task = TriggerDagRunOperator( task_id="test_task", @@ -153,7 +155,13 @@ class TestDagRunOperator: dr = dag_maker.create_dagrun(run_id="test_run_id") ti = dr.get_task_instance(task_id=task.task_id) - mock_xcom_get_one.return_value = ti.run_id + # Mock XCom.get_value to return None for dag_id but return run_id for XCOM_RUN_ID + def mock_get_value(ti_key, key): + if key == XCOM_RUN_ID: + return "test_run_id" + return None + + mock_xcom_get_value.side_effect = mock_get_value link = task.operator_extra_links[0].get_link(operator=task, ti_key=ti.key) @@ -161,6 +169,72 @@ class TestDagRunOperator: expected_url = f"{base_url}dags/{TRIGGERED_DAG_ID}/runs/test_run_id" assert link == expected_url, f"Expected {expected_url}, but got {link}" + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") + @mock.patch("airflow.providers.standard.operators.trigger_dagrun.XCom.get_value") + def test_extra_operator_link_with_dynamic_dag_id(self, mock_xcom_get_value, dag_maker): + """Test that operator link works correctly when dag_id is dynamically resolved from XCom.""" + from airflow.providers.standard.operators.trigger_dagrun import XCOM_DAG_ID, XCOM_RUN_ID + + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): + task = TriggerDagRunOperator( + task_id="test_task", + # In real scenario, this would be a template like "{{ ti.xcom_pull(...) }}" + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="test_run_id", + ) + + dr = dag_maker.create_dagrun(run_id="test_run_id") + ti = dr.get_task_instance(task_id=task.task_id) + + # Mock XCom.get_value to return our test values + def mock_get_value(ti_key, key): + if key == XCOM_DAG_ID: + return "dynamic_dag_id" + if key == XCOM_RUN_ID: + return "dynamic_run_id" + return None + + mock_xcom_get_value.side_effect = mock_get_value + + link = task.operator_extra_links[0].get_link(operator=task, ti_key=ti.key) + + base_url = conf.get("api", "base_url", fallback="/").lower() + # Should use the dag_id from XCom, not the operator attribute + expected_url = f"{base_url}dags/dynamic_dag_id/runs/dynamic_run_id" + assert link == expected_url, f"Expected {expected_url}, but got {link}" + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") + def test_trigger_dagrun_pushes_dag_id_to_xcom(self, dag_maker): + """Test that TriggerDagRunOperator pushes the resolved dag_id to XCom during execution.""" + from airflow.providers.standard.operators.trigger_dagrun import XCOM_DAG_ID + + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + ) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance(task_id=task.task_id) + + # Create a mock task instance that stores XCom values + xcom_values = {} + + def mock_xcom_push(key, value, **kwargs): + xcom_values[key] = value + + ti.xcom_push = mock_xcom_push + + # Execute the task (will raise exception in AF3, but should push XCom first) + try: + task.execute(context={"task_instance": ti}) + except DagRunTriggerException: + pass # Expected in Airflow 3 + + # Verify that the dag_id was pushed to XCom + assert XCOM_DAG_ID in xcom_values + assert xcom_values[XCOM_DAG_ID] == TRIGGERED_DAG_ID + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") def test_trigger_dagrun_custom_run_id(self): task = TriggerDagRunOperator( @@ -577,8 +651,37 @@ class TestDagRunOperatorAF2: assert task.trigger_run_id == "test_run_id" - def test_extra_operator_link(self, dag_maker, session): + def test_trigger_dagrun_pushes_dag_id_to_xcom(self, dag_maker, session): + """Test that TriggerDagRunOperator pushes the resolved dag_id to XCom during execution.""" + from airflow.providers.standard.operators.trigger_dagrun import XCOM_DAG_ID, XCOM_RUN_ID + + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="test_run_id", + ) + dag_maker.create_dagrun() + task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) + + triggering_ti = session.scalar( + select(TaskInstance).filter_by(task_id=task.task_id, dag_id=task.dag_id) + ) + assert triggering_ti is not None + + # Verify that the dag_id was pushed to XCom + dag_id_xcom = triggering_ti.xcom_pull(key=XCOM_DAG_ID) + assert dag_id_xcom == TRIGGERED_DAG_ID + + # Also verify run_id is still pushed + run_id_xcom = triggering_ti.xcom_pull(key=XCOM_RUN_ID) + assert run_id_xcom == "test_run_id" + + @mock.patch("airflow.providers.standard.operators.trigger_dagrun.XCom.get_value") + def test_extra_operator_link(self, mock_xcom_get_value, dag_maker, session): """Asserts whether the correct extra links url will be created.""" + from airflow.providers.standard.operators.trigger_dagrun import XCOM_RUN_ID + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): task = TriggerDagRunOperator( task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, trigger_run_id="test_run_id" @@ -587,13 +690,18 @@ class TestDagRunOperatorAF2: task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) triggering_ti = session.scalar( - select(TaskInstance).where( - TaskInstance.task_id == task.task_id, TaskInstance.dag_id == task.dag_id - ) + select(TaskInstance).filter_by(task_id=task.task_id, dag_id=task.dag_id) ) + # Mock XCom.get_value to return None for dag_id but return run_id for XCOM_RUN_ID + def mock_get_value(ti_key, key): + if key == XCOM_RUN_ID: + return "test_run_id" + return None + + mock_xcom_get_value.side_effect = mock_get_value + with mock.patch("airflow.utils.helpers.build_airflow_url_with_query") as mock_build_url: - # This is equivalent of a task run calling this and pushing to xcom task.operator_extra_links[0].get_link(operator=task, ti_key=triggering_ti.key) assert mock_build_url.called args, _ = mock_build_url.call_args @@ -603,6 +711,47 @@ class TestDagRunOperatorAF2: } assert expected_args in args + @mock.patch("airflow.providers.standard.operators.trigger_dagrun.XCom.get_value") + def test_extra_operator_link_with_dynamic_dag_id(self, mock_xcom_get_value, dag_maker, session): + """Test that operator link works correctly when dag_id is dynamically resolved from XCom.""" + from airflow.providers.standard.operators.trigger_dagrun import XCOM_DAG_ID, XCOM_RUN_ID + + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): + task = TriggerDagRunOperator( + task_id="test_task", + # In real scenario, this would be a template like "{{ ti.xcom_pull(...) }}" + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="test_run_id", + ) + dag_maker.create_dagrun() + task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) + + triggering_ti = session.scalar( + select(TaskInstance).filter_by(task_id=task.task_id, dag_id=task.dag_id) + ) + assert triggering_ti is not None + + # Mock XCom.get_value to return our test values + def mock_get_value(ti_key, key): + if key == XCOM_DAG_ID: + return "dynamic_dag_id" + if key == XCOM_RUN_ID: + return "test_run_id" + return None + + mock_xcom_get_value.side_effect = mock_get_value + + with mock.patch("airflow.utils.helpers.build_airflow_url_with_query") as mock_build_url: + task.operator_extra_links[0].get_link(operator=task, ti_key=triggering_ti.key) + assert mock_build_url.called + args, _ = mock_build_url.call_args + # Should use the dag_id from XCom, not the operator attribute + expected_args = { + "dag_id": "dynamic_dag_id", + "dag_run_id": "test_run_id", + } + assert expected_args in args + def test_trigger_dagrun_with_logical_date(self, dag_maker): """Test TriggerDagRunOperator with custom logical_date.""" custom_logical_date = timezone.datetime(2021, 1, 2, 3, 4, 5) 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 fc832e3..5574e30 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 @@ -4044,7 +4044,17 @@ class TestTriggerDagRunOperator: expected_calls = [ mock.call.send( - msg=TriggerDagRun( + SetXCom( + key="trigger_dag_id", + value="test_dag", + dag_id="test_handle_trigger_dag_run", + task_id="test_task", + run_id="test_run", + map_index=-1, + ), + ), + mock.call.send( + TriggerDagRun( dag_id="test_dag", run_id="test_run_id", reset_dag_run=False, @@ -4052,7 +4062,7 @@ class TestTriggerDagRunOperator: ), ), mock.call.send( - msg=SetXCom( + SetXCom( key="trigger_run_id", value="test_run_id", dag_id="test_handle_trigger_dag_run", @@ -4166,38 +4176,47 @@ class TestTriggerDagRunOperator: assert state == expected_task_state assert msg.state == expected_task_state - expected_calls = [ - mock.call.send( - msg=TriggerDagRun( - dag_id="test_dag", - run_id="test_run_id", - logical_date=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc), - ), - ), - mock.call.send( - msg=SetXCom( - key="trigger_run_id", - value="test_run_id", - dag_id="test_handle_trigger_dag_run_wait_for_completion", - task_id="test_task", - run_id="test_run", - map_index=-1, - ), + # Verify the expected calls were made (order may vary due to SetRenderedFields) + # Check each expected call individually since SetRenderedFields appears first + mock_supervisor_comms.send.assert_any_call( + SetXCom( + key="trigger_dag_id", + value="test_dag", + dag_id="test_handle_trigger_dag_run_wait_for_completion", + task_id="test_task", + run_id="test_run", + map_index=-1, ), - mock.call.send( - msg=GetDagRunState( - dag_id="test_dag", - run_id="test_run_id", - ), + ) + mock_supervisor_comms.send.assert_any_call( + TriggerDagRun( + dag_id="test_dag", + run_id="test_run_id", + logical_date=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc), ), - mock.call.send( - msg=GetDagRunState( - dag_id="test_dag", - run_id="test_run_id", - ), + ) + mock_supervisor_comms.send.assert_any_call( + SetXCom( + key="trigger_run_id", + value="test_run_id", + dag_id="test_handle_trigger_dag_run_wait_for_completion", + task_id="test_task", + run_id="test_run", + map_index=-1, ), + ) + # Verify GetDagRunState was called at least once (may be called multiple times during polling) + get_dag_run_state_calls = [ + call_args + for call_args in mock_supervisor_comms.send.call_args_list + if len(call_args.args) > 0 + and isinstance(call_args.args[0], GetDagRunState) + and call_args.args[0].dag_id == "test_dag" + and call_args.args[0].run_id == "test_run_id" ] - mock_supervisor_comms.assert_has_calls(expected_calls) + assert len(get_dag_run_state_calls) >= 1, ( + f"Expected at least 1 GetDagRunState call, got {len(get_dag_run_state_calls)}" + ) @pytest.mark.parametrize( ("allowed_states", "failed_states", "intermediate_state"),
1 parent 352feb2 commit c46395b

4 files changed

Lines changed: 243 additions & 55 deletions

File tree

‎airflow-core/src/airflow/serialization/serialized_objects.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1587,6 +1587,11 @@ def _deserialize_field_value(cls, field_name: str, value: Any) -> Any:
15871587
elif field_name == "resources":
15881588
return Resources.from_dict(value) if value is not None else None
15891589
elif field_name.endswith("_date"):
1590+
# Check if value is ARG_NOT_SET before trying to deserialize as datetime
1591+
if isinstance(value, dict) and value.get(Encoding.TYPE) == DAT.ARG_NOT_SET:
1592+
from airflow.serialization.definitions.notset import NOTSET
1593+
1594+
return NOTSET
15901595
return cls._deserialize_datetime(value) if value is not None else None
15911596
else:
15921597
# For all other fields, return as-is (strings, ints, bools, etc.)

‎providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py‎

Lines changed: 32 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454

5555
XCOM_LOGICAL_DATE_ISO = "trigger_logical_date_iso"
5656
XCOM_RUN_ID = "trigger_run_id"
57+
XCOM_DAG_ID = "trigger_dag_id"
5758

5859

5960
if TYPE_CHECKING:
@@ -86,21 +87,26 @@ def get_link(self, operator: BaseOperator, *, ti_key: TaskInstanceKey) -> str:
8687
if TYPE_CHECKING:
8788
assert isinstance(operator, TriggerDagRunOperator)
8889

89-
trigger_dag_id = operator.trigger_dag_id
90-
if not AIRFLOW_V_3_0_PLUS:
91-
from airflow.models.renderedtifields import RenderedTaskInstanceFields
92-
from airflow.models.taskinstancekey import TaskInstanceKey as CoreTaskInstanceKey
93-
94-
core_ti_key = CoreTaskInstanceKey(
95-
dag_id=ti_key.dag_id,
96-
task_id=ti_key.task_id,
97-
run_id=ti_key.run_id,
98-
try_number=ti_key.try_number,
99-
map_index=ti_key.map_index,
100-
)
90+
# Try to get the resolved dag_id from XCom first (for dynamic dag_ids)
91+
trigger_dag_id = XCom.get_value(ti_key=ti_key, key=XCOM_DAG_ID)
92+
93+
# Fallback to operator attribute and rendered fields if not in XCom
94+
if not trigger_dag_id:
95+
trigger_dag_id = operator.trigger_dag_id
96+
if not AIRFLOW_V_3_0_PLUS:
97+
from airflow.models.renderedtifields import RenderedTaskInstanceFields
98+
from airflow.models.taskinstancekey import TaskInstanceKey as CoreTaskInstanceKey
99+
100+
core_ti_key = CoreTaskInstanceKey(
101+
dag_id=ti_key.dag_id,
102+
task_id=ti_key.task_id,
103+
run_id=ti_key.run_id,
104+
try_number=ti_key.try_number,
105+
map_index=ti_key.map_index,
106+
)
101107

102-
if template_fields := RenderedTaskInstanceFields.get_templated_fields(core_ti_key):
103-
trigger_dag_id: str = template_fields.get("trigger_dag_id", operator.trigger_dag_id) # type: ignore[no-redef]
108+
if template_fields := RenderedTaskInstanceFields.get_templated_fields(core_ti_key):
109+
trigger_dag_id: str = template_fields.get("trigger_dag_id", operator.trigger_dag_id) # type: ignore[no-redef]
104110

105111
# Fetch the correct dag_run_id for the triggerED dag which is
106112
# stored in xcom during execution of the triggerING task.
@@ -206,7 +212,7 @@ def __init__(
206212
self.note = note
207213
self.deferrable = deferrable
208214
self.logical_date = logical_date
209-
if logical_date is NOTSET:
215+
if isinstance(logical_date, ArgNotSet) or logical_date is NOTSET:
210216
self.logical_date = NOTSET
211217
elif logical_date is None or isinstance(logical_date, (str, datetime.datetime)):
212218
self.logical_date = logical_date
@@ -219,7 +225,7 @@ def __init__(
219225
raise NotImplementedError("Setting `fail_when_dag_is_paused` not yet supported for Airflow 3.x")
220226

221227
def execute(self, context: Context):
222-
if self.logical_date is NOTSET:
228+
if isinstance(self.logical_date, ArgNotSet) or self.logical_date is NOTSET:
223229
# If no logical_date is provided we will set utcnow()
224230
parsed_logical_date = timezone.utcnow()
225231
elif self.logical_date is None or isinstance(self.logical_date, datetime.datetime):
@@ -277,6 +283,14 @@ def execute(self, context: Context):
277283
def _trigger_dag_af_3(self, context, run_id, parsed_logical_date):
278284
from airflow.providers.common.compat.sdk import DagRunTriggerException
279285

286+
# Store the resolved dag_id to XCom for use in the link generation
287+
# This is important for dynamic dag_ids (from XCom or complex templates)
288+
# In Airflow 3.x, context has both "task_instance" and "ti" keys
289+
if "task_instance" in context:
290+
context["task_instance"].xcom_push(key=XCOM_DAG_ID, value=self.trigger_dag_id)
291+
elif "ti" in context:
292+
context["ti"].xcom_push(key=XCOM_DAG_ID, value=self.trigger_dag_id)
293+
280294
kwargs_accepted = dict(
281295
trigger_dag_id=self.trigger_dag_id,
282296
dag_run_id=run_id,
@@ -330,10 +344,11 @@ def _trigger_dag_af_2(self, context, run_id, parsed_logical_date):
330344
raise e
331345
if dag_run is None:
332346
raise RuntimeError("The dag_run should be set here!")
333-
# Store the run id from the dag run (either created or found above) to
347+
# Store the run id and dag_id from the dag run (either created or found above) to
334348
# be used when creating the extra link on the webserver.
335349
ti = context["task_instance"]
336350
ti.xcom_push(key=XCOM_RUN_ID, value=dag_run.run_id)
351+
ti.xcom_push(key=XCOM_DAG_ID, value=self.trigger_dag_id)
337352

338353
if self.wait_for_completion:
339354
# Kick off the deferral process

‎providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py‎

Lines changed: 157 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -140,8 +140,10 @@ def test_trigger_dagrun(self):
140140
assert task.trigger_run_id == expected_run_id # run_id is saved as attribute
141141

142142
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3")
143-
@mock.patch(f"{TRIGGER_OP_PATH}.XCom.get_one")
144-
def test_extra_operator_link(self, mock_xcom_get_one, dag_maker):
143+
@mock.patch("airflow.providers.standard.operators.trigger_dagrun.XCom.get_value")
144+
def test_extra_operator_link(self, mock_xcom_get_value, dag_maker):
145+
from airflow.providers.standard.operators.trigger_dagrun import XCOM_RUN_ID
146+
145147
with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
146148
task = TriggerDagRunOperator(
147149
task_id="test_task",
@@ -153,14 +155,86 @@ def test_extra_operator_link(self, mock_xcom_get_one, dag_maker):
153155
dr = dag_maker.create_dagrun(run_id="test_run_id")
154156
ti = dr.get_task_instance(task_id=task.task_id)
155157

156-
mock_xcom_get_one.return_value = ti.run_id
158+
# Mock XCom.get_value to return None for dag_id but return run_id for XCOM_RUN_ID
159+
def mock_get_value(ti_key, key):
160+
if key == XCOM_RUN_ID:
161+
return "test_run_id"
162+
return None
163+
164+
mock_xcom_get_value.side_effect = mock_get_value
157165

158166
link = task.operator_extra_links[0].get_link(operator=task, ti_key=ti.key)
159167

160168
base_url = conf.get("api", "base_url", fallback="/").lower()
161169
expected_url = f"{base_url}dags/{TRIGGERED_DAG_ID}/runs/test_run_id"
162170
assert link == expected_url, f"Expected {expected_url}, but got {link}"
163171

172+
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3")
173+
@mock.patch("airflow.providers.standard.operators.trigger_dagrun.XCom.get_value")
174+
def test_extra_operator_link_with_dynamic_dag_id(self, mock_xcom_get_value, dag_maker):
175+
"""Test that operator link works correctly when dag_id is dynamically resolved from XCom."""
176+
from airflow.providers.standard.operators.trigger_dagrun import XCOM_DAG_ID, XCOM_RUN_ID
177+
178+
with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
179+
task = TriggerDagRunOperator(
180+
task_id="test_task",
181+
# In real scenario, this would be a template like "{{ ti.xcom_pull(...) }}"
182+
trigger_dag_id=TRIGGERED_DAG_ID,
183+
trigger_run_id="test_run_id",
184+
)
185+
186+
dr = dag_maker.create_dagrun(run_id="test_run_id")
187+
ti = dr.get_task_instance(task_id=task.task_id)
188+
189+
# Mock XCom.get_value to return our test values
190+
def mock_get_value(ti_key, key):
191+
if key == XCOM_DAG_ID:
192+
return "dynamic_dag_id"
193+
if key == XCOM_RUN_ID:
194+
return "dynamic_run_id"
195+
return None
196+
197+
mock_xcom_get_value.side_effect = mock_get_value
198+
199+
link = task.operator_extra_links[0].get_link(operator=task, ti_key=ti.key)
200+
201+
base_url = conf.get("api", "base_url", fallback="/").lower()
202+
# Should use the dag_id from XCom, not the operator attribute
203+
expected_url = f"{base_url}dags/dynamic_dag_id/runs/dynamic_run_id"
204+
assert link == expected_url, f"Expected {expected_url}, but got {link}"
205+
206+
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3")
207+
def test_trigger_dagrun_pushes_dag_id_to_xcom(self, dag_maker):
208+
"""Test that TriggerDagRunOperator pushes the resolved dag_id to XCom during execution."""
209+
from airflow.providers.standard.operators.trigger_dagrun import XCOM_DAG_ID
210+
211+
with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
212+
task = TriggerDagRunOperator(
213+
task_id="test_task",
214+
trigger_dag_id=TRIGGERED_DAG_ID,
215+
)
216+
217+
dr = dag_maker.create_dagrun()
218+
ti = dr.get_task_instance(task_id=task.task_id)
219+
220+
# Create a mock task instance that stores XCom values
221+
xcom_values = {}
222+
223+
def mock_xcom_push(key, value, **kwargs):
224+
xcom_values[key] = value
225+
226+
ti.xcom_push = mock_xcom_push
227+
228+
# Execute the task (will raise exception in AF3, but should push XCom first)
229+
try:
230+
task.execute(context={"task_instance": ti})
231+
except DagRunTriggerException:
232+
pass # Expected in Airflow 3
233+
234+
# Verify that the dag_id was pushed to XCom
235+
assert XCOM_DAG_ID in xcom_values
236+
assert xcom_values[XCOM_DAG_ID] == TRIGGERED_DAG_ID
237+
164238
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3")
165239
def test_trigger_dagrun_custom_run_id(self):
166240
task = TriggerDagRunOperator(
@@ -583,8 +657,37 @@ def test_explicitly_provided_trigger_run_id_is_saved_as_attr(self, dag_maker, se
583657

584658
assert task.trigger_run_id == "test_run_id"
585659

586-
def test_extra_operator_link(self, dag_maker, session):
660+
def test_trigger_dagrun_pushes_dag_id_to_xcom(self, dag_maker, session):
661+
"""Test that TriggerDagRunOperator pushes the resolved dag_id to XCom during execution."""
662+
from airflow.providers.standard.operators.trigger_dagrun import XCOM_DAG_ID, XCOM_RUN_ID
663+
664+
with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
665+
task = TriggerDagRunOperator(
666+
task_id="test_task",
667+
trigger_dag_id=TRIGGERED_DAG_ID,
668+
trigger_run_id="test_run_id",
669+
)
670+
dag_maker.create_dagrun()
671+
task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
672+
673+
triggering_ti = session.scalar(
674+
select(TaskInstance).filter_by(task_id=task.task_id, dag_id=task.dag_id)
675+
)
676+
assert triggering_ti is not None
677+
678+
# Verify that the dag_id was pushed to XCom
679+
dag_id_xcom = triggering_ti.xcom_pull(key=XCOM_DAG_ID)
680+
assert dag_id_xcom == TRIGGERED_DAG_ID
681+
682+
# Also verify run_id is still pushed
683+
run_id_xcom = triggering_ti.xcom_pull(key=XCOM_RUN_ID)
684+
assert run_id_xcom == "test_run_id"
685+
686+
@mock.patch("airflow.providers.standard.operators.trigger_dagrun.XCom.get_value")
687+
def test_extra_operator_link(self, mock_xcom_get_value, dag_maker, session):
587688
"""Asserts whether the correct extra links url will be created."""
689+
from airflow.providers.standard.operators.trigger_dagrun import XCOM_RUN_ID
690+
588691
with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
589692
task = TriggerDagRunOperator(
590693
task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, trigger_run_id="test_run_id"
@@ -593,13 +696,18 @@ def test_extra_operator_link(self, dag_maker, session):
593696
task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
594697

595698
triggering_ti = session.scalar(
596-
select(TaskInstance).where(
597-
TaskInstance.task_id == task.task_id, TaskInstance.dag_id == task.dag_id
598-
)
699+
select(TaskInstance).filter_by(task_id=task.task_id, dag_id=task.dag_id)
599700
)
600701

702+
# Mock XCom.get_value to return None for dag_id but return run_id for XCOM_RUN_ID
703+
def mock_get_value(ti_key, key):
704+
if key == XCOM_RUN_ID:
705+
return "test_run_id"
706+
return None
707+
708+
mock_xcom_get_value.side_effect = mock_get_value
709+
601710
with mock.patch("airflow.utils.helpers.build_airflow_url_with_query") as mock_build_url:
602-
# This is equivalent of a task run calling this and pushing to xcom
603711
task.operator_extra_links[0].get_link(operator=task, ti_key=triggering_ti.key)
604712
assert mock_build_url.called
605713
args, _ = mock_build_url.call_args
@@ -609,6 +717,47 @@ def test_extra_operator_link(self, dag_maker, session):
609717
}
610718
assert expected_args in args
611719

720+
@mock.patch("airflow.providers.standard.operators.trigger_dagrun.XCom.get_value")
721+
def test_extra_operator_link_with_dynamic_dag_id(self, mock_xcom_get_value, dag_maker, session):
722+
"""Test that operator link works correctly when dag_id is dynamically resolved from XCom."""
723+
from airflow.providers.standard.operators.trigger_dagrun import XCOM_DAG_ID, XCOM_RUN_ID
724+
725+
with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
726+
task = TriggerDagRunOperator(
727+
task_id="test_task",
728+
# In real scenario, this would be a template like "{{ ti.xcom_pull(...) }}"
729+
trigger_dag_id=TRIGGERED_DAG_ID,
730+
trigger_run_id="test_run_id",
731+
)
732+
dag_maker.create_dagrun()
733+
task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
734+
735+
triggering_ti = session.scalar(
736+
select(TaskInstance).filter_by(task_id=task.task_id, dag_id=task.dag_id)
737+
)
738+
assert triggering_ti is not None
739+
740+
# Mock XCom.get_value to return our test values
741+
def mock_get_value(ti_key, key):
742+
if key == XCOM_DAG_ID:
743+
return "dynamic_dag_id"
744+
if key == XCOM_RUN_ID:
745+
return "test_run_id"
746+
return None
747+
748+
mock_xcom_get_value.side_effect = mock_get_value
749+
750+
with mock.patch("airflow.utils.helpers.build_airflow_url_with_query") as mock_build_url:
751+
task.operator_extra_links[0].get_link(operator=task, ti_key=triggering_ti.key)
752+
assert mock_build_url.called
753+
args, _ = mock_build_url.call_args
754+
# Should use the dag_id from XCom, not the operator attribute
755+
expected_args = {
756+
"dag_id": "dynamic_dag_id",
757+
"dag_run_id": "test_run_id",
758+
}
759+
assert expected_args in args
760+
612761
def test_trigger_dagrun_with_logical_date(self, dag_maker):
613762
"""Test TriggerDagRunOperator with custom logical_date."""
614763
custom_logical_date = timezone.datetime(2021, 1, 2, 3, 4, 5)

0 commit comments

Comments
 (0)