Module 28: Run the Project with Airflow

export AIRFLOW_CONN_DATABRICKS_DEFAULT='databricks://@dbc-xxxx.cloud.databricks.com?token=<url-encoded-token>'

pip install apache-airflow-providers-databricks
# optional: Azure identity support
pip install "apache-airflow-providers-databricks[azure-identity]"

# dags/check_databricks.py
from datetime import datetime
from airflow import DAG
from airflow.decorators import task
from airflow.providers.databricks.hooks.databricks import DatabricksHook

with DAG("check_databricks", start_date=datetime(2025, 1, 1), schedule=None, catchup=False):

    @task
    def ping():
        hook = DatabricksHook(databricks_conn_id="databricks_default")
        for job in hook.list_jobs(limit=5):          # a harmless read-only call
            print(job.get("settings", {}).get("name"))

    ping()

# dags/orders_pipeline.py
from datetime import datetime, timedelta
from airflow import DAG
from airflow.providers.databricks.operators.databricks import DatabricksRunNowOperator

with DAG(
    dag_id="orders_pipeline",
    start_date=datetime(2025, 1, 1),
    schedule="30 6 * * *",
    catchup=False,
    default_args={"retries": 1, "retry_delay": timedelta(minutes=5)},
) as dag:
    run_job = DatabricksRunNowOperator(
        task_id="run_orders_job",
        databricks_conn_id="databricks_default",
        job_name="orders-pipeline-prod",            # or job_id=123456789
        job_parameters={"run_date": "{{ ds }}"},
    )

from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator

one_off = DatabricksSubmitRunOperator(
    task_id="adhoc_gold_rebuild",
    run_name="orders-gold-rebuild-{{ ds }}",
    tasks=[{
        "task_key": "build_gold",
        "python_wheel_task": {
            "package_name": "orders_pipeline",
            "entry_point": "build_gold",
            "named_parameters": {"catalog": "<catalog>", "run_date": "{{ ds }}", "days": "30"},
        },
        "environment_key": "default",
    }],
    environments=[{
        "environment_key": "default",
        "spec": {"environment_version": "3",
                 "dependencies": ["/Volumes/<catalog>/dbx_project/artifacts/orders_pipeline-0.1.0-py3-none-any.whl"]},
    }],
)

from airflow.providers.databricks.operators.databricks_workflow import DatabricksWorkflowTaskGroup
from airflow.providers.databricks.operators.databricks import (
    DatabricksNotebookOperator, DatabricksTaskOperator,
)

with DAG("orders_workflow", start_date=datetime(2025, 1, 1), schedule=None, catchup=False) as dag:
    group = DatabricksWorkflowTaskGroup(
        group_id="orders",
        databricks_conn_id="databricks_default",
        job_clusters=job_cluster_spec,                 # shared cluster definition
        notebook_params={"run_date": "{{ ds }}"},      # put parameters here, not on each task
    )
    with group:
        ingest = DatabricksNotebookOperator(
            task_id="ingest", databricks_conn_id="databricks_default",
            notebook_path="/Shared/ingest", source="WORKSPACE",
            job_cluster_key="Shared_job_cluster",
        )
        report = DatabricksTaskOperator(
            task_id="report", databricks_conn_id="databricks_default",
            task_config={"sql_task": {"query": {"query_id": "<QUERY_ID>"},
                                      "warehouse_id": "<WAREHOUSE_ID>"}},
        )
        ingest >> report

run_job = DatabricksRunNowOperator(
    task_id="run_orders_job",
    job_name="orders-pipeline-prod",
    job_parameters={
        "run_date": "{{ ds }}",
        "days": "{{ var.value.get('orders_days', '7') }}",
    },
)

from airflow.decorators import task

@task
def log_run(ti=None):
    run_id = ti.xcom_pull(task_ids="run_orders_job", key="run_id")
    url = ti.xcom_pull(task_ids="run_orders_job", key="run_page_url")
    print(f"databricks run_id={run_id} url={url}")

run_job >> log_run()

run_job = DatabricksRunNowOperator(
    task_id="run_orders_job", job_name="orders-pipeline-prod",
    job_parameters={"run_date": "{{ ds }}"},
    execution_timeout=timedelta(hours=1),
    retries=1, repair_run=True,
)

# databricks.yml, prod target: Airflow owns the trigger, so no schedule block
resources:
  jobs:
    orders_pipeline:
      name: orders-pipeline-${bundle.target}
      # schedule: removed in prod