Using return value from `TriggerDagRunOperator` in parent dag (pulling `xcom` from a given `run_id`)

[hterik]

Description

When triggering a dag through TriggerDagRunOperator, it would be useful for the dag that initiated the trigger to fetch information produced in the triggered dag .

Example:

@dag()
def child_dag():
    @task
    def run():
        return "hello world"
    run()

@dag()
def top_dag():
    @task
    def postprocess(trigger_task):
        context = get_current_context()
        child_run_id = context["task_instance"].xcom_pull(task_ids="trigger_child", key="trigger_run_id")
        child_dag_return_value = context["task_instance"].xcom_pull(
            dag_id="child_dag",
            task_id="run",
            DESIRED_PARAMETER_TO_SET_RUN_ID=child_run_id       # <---------------------------------------
        )
        assert child_dag_return_value == "hello world"


    t_trigger = TriggerDagRunOperator(
        task_id="trigger_child",
        trigger_dag_id="child_dag",
        wait_for_completion=True
    )
    postprocess(t_trigger)

While xcom_pull does allow fetching from a different dag_id, it has no parameter to control the exact run_id, it assumes the execution_date is the same, or with include_prior_dates=True it takes the latest found. (In the example above this doesn't make any difference because the return value is static, in more realistic scenarios, the output is based on parameters passed from the top_dag and need to be in sync)

Use case/motivation

No response

Related issues

No response

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[jroachgolf84]

So, a couple of things here. XComs are tied to a TaskInstance, not a DagRun. That is, having a snippet of code like the one below is ambiguous, as there are multiple Tasks that an XCom could be written for.

    t_trigger = TriggerDagRunOperator(
        task_id="trigger_child",
        trigger_dag_id="child_dag",
        wait_for_completion=True
    )
    postprocess(t_trigger)

Also, there is a run_id parameter that can be passed to xcom_pull (see the code here: https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/models/taskinstance.py#L1671-L1699). I think a docs improvement might be applicable here, but IMO, the functionality is already in-place.

[tirkarthi]

Agree with @jroachgolf84 that this could be a documentation improvement. Following dag works as expected.

from __future__ import annotations

from airflow.providers.standard.operators.python import get_current_context
from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.sdk import dag, task


@dag()
def gh62988_child_dag():
    @task
    def run():
        return "hello world"

    run()


@dag()
def gh62988_top_dag():
    @task
    def postprocess(trigger_task):
        context = get_current_context()
        print(context["task_instance"])
        child_run_id = context["task_instance"].xcom_pull(task_ids="trigger_dag", key="trigger_run_id")
        print(f"{child_run_id = }")
        child_dag_return_value = context["task_instance"].xcom_pull(
            dag_id="gh62988_child_dag", task_ids="run", run_id=child_run_id
        )
        print(f"{child_dag_return_value = }")
        if not child_dag_return_value == "hello world":
            raise ValueError(f"child dag return_value {child_dag_return_value}")

    t_trigger = TriggerDagRunOperator(
        task_id="trigger_dag", trigger_dag_id="gh62988_child_dag", wait_for_completion=True
    )
    postprocess(t_trigger.output)


gh62988_top_dag()
gh62988_child_dag()

[jroachgolf84]

This PR has already been linked to the issue, but I think that it does a good job adding documentation and unit-tests: #63030 (review)