CODEWITHSUNDEEP
airflow
Kay
Kay

Reputation: 2332

Using dag_run variables in airflow Dag

I am trying to use airflow variables to determine whether to execute a task or not. I have tried this and it's not working:

if '{{ params.year }}' == '{{ params.message }}':
     run_this = DummyOperator (
                task_id = 'dummy_dag'
               )

I was hoping to get some help making it work. Also is there a better way of doing something like this in airflow?

Upvotes: 0

Views: 4170

Answers (1)

NicoE
NicoE

Reputation: 4853

I think a good way to solve this, is with BranchPythonOperator to branch dynamically based on the provided DAG parameters. Consider this example:

Use params to provide the parameters to the DAG (could be also done from the UI), in this example: {"enabled": True}

from airflow.decorators import dag, task
from airflow.utils.dates import days_ago
from airflow.operators.python import get_current_context, BranchPythonOperator

@dag(
    default_args=default_args,
    schedule_interval=None,
    start_date=days_ago(1),
    catchup=False,
    tags=["example"],
    params={"enabled": True},
)
def branch_from_dag_params():
    def _print_enabled():
        context = get_current_context()
        enabled = context["params"].get("enabled", False)
        print(f"Task id: {context['ti'].task_id}")
        print(f"Enabled is: {enabled}")

    @task
    def task_a():
        _print_enabled()

    @task
    def task_b():
        _print_enabled()

Define a callable to the BranchPythonOperator in which you will perform your conditionals and return the next task to be executed. You can access the execution context variables from **kwargs. Also keep in mind that this operator should return a single task_id or a list of task_ids to follow downstream. Those resultant tasks should always be directly downstream from it.

    def _get_task_run(ti, **kwargs):
        custom_param = kwargs["params"].get("enabled", False)

        if custom_param:
            return "task_a"
        else:
            return "task_b"

    branch_task = BranchPythonOperator(
        task_id="branch_task",
        python_callable=_get_task_run,
    )
    task_a_exec = task_a()
    task_b_exec = task_b()
    branch_task >> [task_a_exec, task_b_exec]

The result is that task_a gets executed and task_b is skipped :

branch from dag params

AIRFLOW_CTX_DAG_OWNER=airflow
AIRFLOW_CTX_DAG_ID=branch_from_dag_params
AIRFLOW_CTX_TASK_ID=task_a
Task id: task_a
Enabled is: True

Let me know if that worked for you.

Docs

Upvotes: 1

Related Questions