Databricks Provider for Apache Airflow

from __future__ import annotations

import os
import pendulum

from airflow.models.dag import DAG
from airflow.providers.databricks.operators.databricks import DatabricksRunNowOperator

# Configure your Databricks connection in Airflow UI (Admin -> Connections)
# Conn Id: 'databricks_default'
# Conn Type: 'Databricks'
# Host: Your Databricks workspace URL (e.g., https://adb-xxxxxxxx.xx.databricks.com)
# Password: Your Databricks Personal Access Token (PAT) (recommended over username/password)
# See https://airflow.apache.org/docs/apache-airflow-providers-databricks/stable/connections/databricks.html

with DAG(
    dag_id="databricks_run_notebook_example",
    start_date=pendulum.datetime(2023, 1, 1, tz="UTC"),
    schedule=None,
    catchup=False,
    tags=["databricks", "example"],
    doc_md="""
    ### Databricks Run Notebook Example DAG
    This DAG demonstrates how to use `DatabricksRunNowOperator` to execute
    an existing Databricks job (which might run a notebook or JAR).

    **Prerequisites:**
    1. An Airflow connection named `databricks_default` configured with your Databricks workspace details and PAT.
    2. An existing Databricks job. Replace `YOUR_DATABRICKS_JOB_ID` with your actual Job ID.
    """,
) as dag:
    run_databricks_job = DatabricksRunNowOperator(
        task_id="run_existing_databricks_job",
        databricks_conn_id="databricks_default",
        job_id="YOUR_DATABRICKS_JOB_ID", # Replace with your Databricks Job ID
        # You can pass optional parameters to the job run
        # notebook_params={"input_param": "value_from_airflow"},
        # or spark_jar_params, spark_python_task_params, etc. depending on job type
    )