airflow access dag_run.conf values and pass it to SSHOperator command

Question:

If I use this json parameter to trigger a DAG {"number":["1","2","3"], "items":["a", "b", "c"]}, how can I use dag_run.conf to access the values and pass them to the ssh command in the SSHoperator task?

For example

def run_config(**kwargs):
    number_input = kwargs['dag_run'].conf.get('number')
    items_input = kwargs['dag_run'].conf.get('items')
    return number_input, items_input

#assign the returned values from the run_config function to the following variables
number, item = run_config

#use the above 2 variables as argument in the following ssh command call 
with DAG(blah blah blah) as dag:

    for i in range(0, len(number)):
        process_task = SSHOperator(
                    task_id='run_'.format(number[i]),
                    command='run_somthing.sh' + ' ' + number[i] + ' ' + item[i]
                    )


Asked By: Dozel

||

Answers:

If you are using a version >= 2.3.0, you can use the Dynamic Task Mapping:

from airflow import DAG
from airflow.decorators import task
from airflow.providers.ssh.operators.ssh import SSHOperator

with DAG(dag_id="...") as dag:
    @task
    def get_config(**context):
        return context['dag_run'].conf.get('items')

    items_to_process = get_config()

    SSHOperator.partial(
        task_id='run_ssh'
    ).expand(
        command=[f"run_something.sh {ind} {val}" for ind, val in enumerate(items_to_process)]
    )

For the versions < 2.3.0, you can inspire from this answer

Answered By: Hussein Awala

actually, I could use the jinja template directly to get the trigger parameter value into any operator without using a function nor pythonOperator to call it.

will mark the answer from @hussein Awala as correct since a dynamic creating tasks method is introduced, for which can’t be done with my solution below.

process_task = SSHOperator(
        task_id='run_job',
        command=shell_command + ' ' + "{{ dag_run.conf['number'] }}" + ' ' + "{{ dag_run.conf['item'] }}"
)
Answered By: Dozel
Categories: questions Tags: ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.