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

Question

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")

Answer 1

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.

Answer 2

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")