Can TriggerDagRunOperator have a conf that is a number derived from a xcom or task?

radiomime 0
Viewed 29

I'm newer to airflow, but I'm having difficulties really understanding how to pass small xcom values around. Sometimes, this seems to work without an issue; other times, it takes me hours. Here's an example that has really been getting me.

I'm trying to trigger an external DAG, and I need to pass a number as a param to the external dag. I don't know this number, but I get it from a task in my dag. How can I pass this as a conf/param to the triggered DAG?

This is probably easiest to understand through an example:

import datetime

from airflow import DAG
from airflow.decorators import task
from airflow.models import Param
from airflow.operators.bash import BashOperator
from airflow.operators.trigger_dagrun import TriggerDagRunOperator

with DAG(
    dag_id="example_trigger_target_dag",
    schedule_interval=None,
    start_date=datetime.datetime(year=2022, month=1, day=1),
    catchup=False,
    tags=["example"],
    params={
        "a_number": Param(
            1,
            type="number",
        ),
    },
) as target_dag:
    bash_task = BashOperator(
        task_id="bash_task",
        bash_command="echo 'Here is some number: {{ params.a_number }}'",
    )


with DAG(
    dag_id="example_trigger_controller_dag",
    schedule_interval=None,
    start_date=datetime.datetime(year=2022, month=1, day=1),
    catchup=False,
    tags=["example"],
    params={
        "different_number": Param(
            2,
            type="number",
        ),
    },
) as trigger_dag:
    trigger_with_2 = TriggerDagRunOperator(
        task_id="test_trigger_dagrun_with_2",
        trigger_dag_id="example_trigger_target_dag",  # Ensure this equals the dag_id of the DAG to trigger
        conf={
            "a_number": 2,
        },
    )

    ##############################
    # I'd expect this to work, but it doesn't.
    # error: during build
    # airflow.exceptions.AirflowException: conf parameter should be JSON Serializable
    ##########

    # @task
    # def get_number_3() -> int:
    #     return 3
    #
    # number_3_from_task = get_number_3()
    # trigger_with_3 = TriggerDagRunOperator(
    #     task_id="test_trigger_dagrun_with_3",
    #     trigger_dag_id="example_trigger_target_dag",  # Ensure this equals the dag_id of the DAG to trigger
    #     conf={
    #         "a_number": number_3_from_task,
    #     },
    # )

    ##############################
    # This was only in response to the above. It feels worse and uglier, but I'd have made do with it. But it, too, doesn't work.
    # error during dag run
    # Failed validating 'type' in schema:
    #     {'type': 'number'}
    ##########
    @task
    def get_number_4() -> int:
        return 4

    number_4_from_task = get_number_4()
    trigger_with_4 = TriggerDagRunOperator(
        task_id="test_trigger_dagrun_with_4",
        trigger_dag_id="example_trigger_target_dag",  # Ensure this equals the dag_id of the DAG to trigger
        wait_for_completion=True,
        conf={
            "a_number": "{{ ti.xcom_pull(task_ids='get_number_4', key='return_value') | int }}",
        },
    )
    number_4_from_task >> trigger_with_4

    ##############################
    # This works without an issue.
    # I've tried wrapping the TriggerDagRunOperator in a decorated task, but I have issues waiting for that task to finish.
    # Also, it doesn't seem to solve the issue I'm having here, anyways.
    ##########
    @task
    def print_a_number(your_num: int) -> str:
        text = f"wow, a number: {your_num}"

        print(text)
        return text

    print_a_number(number_4_from_task)
airflow-2.x airflow python
1 Answers

I found a solution, but it isn't the most intuitive way for this to work.

It also produces new behavior for Operators within the Triggering DAG. Any DAG that may need to trigger another DAG might run into issues with passing env variables to a BashOperator, for instance. Still, in absence of a better solution, this is okay; the BashOperator env problem could be solved with Jinja Templating, for instance.

import datetime

from airflow import DAG
from airflow.decorators import task
from airflow.models import Param
from airflow.operators.bash import BashOperator
from airflow.operators.trigger_dagrun import TriggerDagRunOperator

with DAG(
    dag_id="example_trigger_target_dag",
    schedule_interval=None,
    start_date=datetime.datetime(year=2022, month=1, day=1),
    catchup=False,
    tags=["example"],
    params={
        "a_number": Param(
            1,
            type="number",
        ),
    },
) as target_dag:
    bash_task = BashOperator(
        task_id="bash_task",
        bash_command="echo 'Here is some number: {{ params.a_number }}'",
    )


with DAG(
    dag_id="example_trigger_controller_dag",
    schedule_interval=None,
    start_date=datetime.datetime(year=2022, month=1, day=1),
    catchup=False,
    # Below is the important bit!
    render_template_as_native_obj=True,
    tags=["example"],
    params={
        "different_number": Param(
            2,
            type="number",
        ),
    },
) as trigger_dag:

    ##############################
    # this now works, thanks to 'render_template_as_native_obj=True'
    # this introduces problems of it's own. Passing env to a BashOperator like:
    # env={**os.environ, FOO=BAR} will now produce errors, as "" is None and "123" becomes 123 for environment variables
    ##########
    @task
    def get_number_4() -> int:
        return 4

    number_4_from_task = get_number_4()
    trigger_with_4 = TriggerDagRunOperator(
        task_id="test_trigger_dagrun_with_4",
        trigger_dag_id="example_trigger_target_dag",  # Ensure this equals the dag_id of the DAG to trigger
        wait_for_completion=True,
        conf={
            "a_number": "{{ ti.xcom_pull(task_ids='get_number_4', key='return_value') | int }}",
        },
    )
    number_4_from_task >> trigger_with_4

    ##############################
    # I think this would be the most intuitive way, but I haven't been able to figure it out
    # error: during build
    # airflow.exceptions.AirflowException: conf parameter should be JSON Serializable
    ##########

    # @task
    # def get_number_3() -> int:
    #     return 3
    #
    # number_3_from_task = get_number_3()
    # trigger_with_3 = TriggerDagRunOperator(
    #     task_id="test_trigger_dagrun_with_3",
    #     trigger_dag_id="example_trigger_target_dag",  # Ensure this equals the dag_id of the DAG to trigger
    #     conf={
    #         "a_number": number_3_from_task,
    #     },
    # )
radiomime0
Related