CODEWITHSUNDEEP
airflowdatabricks
Eric Bellet
Eric Bellet

Reputation: 2045

How to use the result of a BashOperator task as argument of another Airflow task?

I need to pass a job_id parameter to my object DatabricksRunNowOperator(). The job_id is the result of executing the databricks jobs create --json '{myjson} command.

$ databricks jobs create --json '{myjson}'

{job_id: 12}

import os
import subprocess    
from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from airflow.contrib.operators.databricks_operator import DatabricksRunNowOperator

def pull_function():
    returned_output = subprocess.check_output("echo ti.xcom_pull(key='jobid_CreateCreateRobot')")
    return returned_output


dag_CreateRobot = DAG(dag_id='CreateRobot', 
                      default_args={'owner': 'eric', 
                                  'email': [],
                                  'depends_on_past': False, 
                                  'start_date':'2019-09-16 16:48:28.803023', 
                                  'provide_context': True}, 
                      schedule_interval='@once')

CreateRobot = BashOperator(dag=dag_CreateRobot, 
                                 task_id='CreateRobot', 
                                 bash_command="databricks jobs create --json '{myjson}')")\

RunRobot = DatabricksRunNowOperator(dag=dag_CreateRobot, 
                                    task_id=ti.xcom_pull('RunCreateRobot'), 
                                    job_id=pull_function(), 
                                    databricks_conn_id='myconn', 
                                    json={'token': 'mytoken' })

RunRobot.set_upstream(CreateRobot)

I wrote this code for explaining my goal but it does not work. How can I do for using the result of a BashOperator task into other task that depends of it?

Upvotes: 2

Views: 4149

Answers (1)

jay.cs
jay.cs

Reputation: 303

The bash command in the BashOperator needs to be $ databricks jobs create --json '{myjson}'

i.e.

CreateRobot = BashOperator(dag=dag_CreateRobot, 
                            task_id='CreateRobot', 
                            bash_command="databricks jobs create --json '{myjson}')",
                            xcom_push=True #Specify this in older airflow versions)

The above operator when executed pushes the last line of the output to xcom. (https://airflow.apache.org/_modules/airflow/operators/bash_operator.html)

The xcom value can be accessed using :

ti.xcom_pull(task_ids='CreateRobot')

Upvotes: 2

Related Questions