airflow_dag_databricks_submit_run_operator_spark_job.py

1import os
2from datetime import datetime
3
4from airflow import DAG
5from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator
6
7# Example of task parameters
8new_cluster = {
9    "spark_version": "14.3.x-scala2.12",
10    "node_type_id": "i3.xlarge",
11    "num_workers": 2,
12}
13
14notebook_task = {
15    "notebook_path": "/Users/airflow@example.com/PrepareData",
16}
17
18with DAG(
19    dag_id="example_databricks_operator",
20    start_date=datetime(2021, 1, 1),
21    schedule_interval=None,
22    catchup=False,
23    tags=["example"],
24) as dag:
25    # [START howto_operator_databricks_submit_run]
26    task_submit_run = DatabricksSubmitRunOperator(
27        task_id="submit_run",
28        databricks_conn_id="databricks_default",
29        new_cluster=new_cluster,
30        notebook_task=notebook_task,
31    )
32    # [END howto_operator_databricks_submit_run]