Can I pass variable values to Python script from PythonOperator in Airflow using op_args or op_kwargs?

Viewed 411

From airflow DAG can I pass variable values as arguments to the python script using op_args or op_kwargs. In my airflow, Dag is importing my script as from scripts import my_script my python operator looks like this

PythonOperator(
    task_id='xxxxxx',
    python_callable=my_script.main,
    op_args=[bucket_name, prefix, source_blob_name, dest_bucket_name],
    dag=dag,
    trigger_rule='all_success'
)

I declared my variables in Airflow. I could be able to call my values bucket_name=Variable.get('bucket_name') here I want pass the value of bucket_name to my variable in Python script is that possible to do so?

2 Answers

In PythonOperator op_args, op_kwargs, templates_dict are templated fields.

So you can do :

PythonOperator(
    ...,
    op_args=['{{ var.value.bucket_name }}'],
    python_callable=my_script.main
)

Then your Python callable would be:

def main(*op_args):
    bucket_name = op_args[0]

You can also use op_kwargs / templates_dict:

PythonOperator(
    ...,
    templates_dict={'bucket_name', '{{ var.value.bucket_name }}'},
    python_callable=my_script.main
)

Then your Python callable would be:

def main(bucket_name, **context):
    ...

But there is no need to do either of them. There is no reason to pass argument that you can access directly with in your Python callable.

You can just do:

from airflow.models.variable import Variable
def main(**context):
    bucket_name = Variable.get('bucket_name')

This is perfectly safe since main is called only when PythonOperator is executed.

You can use params inside your PythonOperator like this:

PythonOperator(
...
params={"bucket_name": bucket_name, "key": value}
...
)

And retrieve it:

def main(**kwargs):
 bucket_name = kwargs.get("bucket_name")
...
Related