How do I use an Airflow variable inside a Databricks notebook?

user_2197662 0
Viewed 854

I have a Databricks PySpark notebook that gets called from an Airflow DAG. I created a variable in Airflow by going to Admin - Variables and added a key-value pair.

I cannot find a way to use that Airflow variable in Databricks.

Edit to add sample of my code.

notebook_task = {
    'notebook_path': '/Users/email@exaple.com/myDAG',
    'base_parameters': {
        "token": token
    }
}

and the operator defined here

opr_submit_run = DatabricksSubmitRunOperator(
    task_id='run_notebook',
    existing_cluster_id='xxxxx', 
    run_name='test',
    databricks_conn_id='databricks_xxx',  
    notebook_task=notebook_task
)

What ended up working is using base_parameters instead of notebook_parans which can be found here https://docs.databricks.com/dev-tools/api/latest/jobs.html

and accessing it from databricks by using

my_param = dbutils.widgets.get("token")
airflow pyspark azure-databricks databricks
2 Answers

if you set it as a parameter to the notebook call (parameters inside notebook_task), then you need to use the dbutils.widgets.get function, put at the beginning of notebook something like this:

my_param = dbutils.widgets.get("key")
Alex_Ott_186270

Extending the answer provided by Alex since this question was asked in the context of Apache-Airflow that executing a databricks notebook.

The DatabricksRunNowOperator (which is available by the databricks provider) has notebook_params that is a dict from keys to values for jobs with notebook task, e.g. "notebook_params": {"name": "john doe", "age": "35"}. The map is passed to the notebook and will be accessible through the dbutils.widgets.get function. As Alex explained you can access the value from databricks notebook as:

my_param = dbutils.widgets.get("key")

An example usage will be:

spark_jar_task = DatabricksSubmitRunOperator(
    task_id='spark_jar_task',
    new_cluster=new_cluster,
    notebook_params={"name": "john doe", "age":  "35"},
    spark_jar_task={'main_class_name': 'com.example.ProcessData'},
    libraries=[{'jar': 'dbfs:/lib/etl-0.1.jar'}],
)

The issue now is how to pass a value from Airflow Variable rather than a static value. For that we need the notebook_params to be a templated field so the Jinja engine will template the value. The problem is that notebook_params is not listed in the template_fields To overcome this we can create a custom version of the operator as:

class MyDatabricksRunNowOperator(DatabricksRunNowOperator):
    template_fields = DatabricksRunNowOperator.template_fields + ('notebook_params',)

Then we can use macro {{ var.value.my_var }} which will be templated during run time as:

spark_jar_task = MyDatabricksSubmitRunOperator(
    task_id='spark_jar_task',
    new_cluster=new_cluster,
    notebook_params={"var_value": {{ var.value.my_var }} },
    spark_jar_task={'main_class_name': 'com.example.ProcessData'},
    libraries=[{'jar': 'dbfs:/lib/etl-0.1.jar'}],
)

The operator will get the value of my_var Variable and pass it to your notebook.

user_146244090
Related