moomindani commented on PR #66613:
URL: https://github.com/apache/airflow/pull/66613#issuecomment-4410735105
## Real-environment validation
Ran the operators end-to-end through `airflow dags test` against a real
Databricks workspace, with verification tasks in the same DAG that read the
workspace state back via the REST API. All seven tasks succeeded.
### Tasks (DAG `pr66613_realenv`, with DAG-level `params={"env": "prod",
"batch_size": "100"}`)
| Task | Status |
|---|---|
| `create_job` (`DatabricksCreateJobsOperator`) | success |
| `verify_create_job` — assert workspace `settings.parameters` contains
`env=prod`, `batch_size=100` | success |
| `run_now` (`DatabricksRunNowOperator`, `params={"env": "staging",
"batch_size": "42"}`) | success |
| `verify_run_now` — assert run's `job_parameters` is `{"env": "staging",
"batch_size": "42"}` (operator overrides DAG) | success |
| `submit_run` (`DatabricksSubmitRunOperator`, `params={"env": "dev",
"shard": "1"}`) | success |
| `verify_submit_run` — assert task's `notebook_task.base_parameters` is
`{"env": "dev", "batch_size": "100", "shard": "1"}` (DAG-level `batch_size` is
inherited where the operator does not override) | success |
| `cleanup_job` | success |
The verification tasks confirm the params actually arrive at Databricks (not
just that the request body is constructed locally), and that DAG-level params
and operator-level params merge correctly.
### DAG
<details>
<summary>dev/dag_pr66613_realenv.py</summary>
```python
"""Real-environment validation DAG for PR #66613 (GH-39002)."""
from __future__ import annotations
import os
from datetime import datetime
from airflow.providers.databricks.hooks.databricks import DatabricksHook
from airflow.providers.databricks.operators.databricks import (
DatabricksCreateJobsOperator,
DatabricksRunNowOperator,
DatabricksSubmitRunOperator,
)
from airflow.sdk import DAG, task
NOTEBOOK_PATH = os.environ.get(
"PR66613_NOTEBOOK_PATH",
"/Users/<your-user>@example.com/airflow-pr66613-noop"
)
def _hook() -> DatabricksHook:
return DatabricksHook(databricks_conn_id="databricks_default")
with DAG(
dag_id="pr66613_realenv",
start_date=datetime(2026, 1, 1),
schedule=None,
catchup=False,
params={"env": "prod", "batch_size": "100"},
tags=["databricks", "pr66613"],
) as dag:
create_job = DatabricksCreateJobsOperator(
task_id="create_job",
json={
"name": "{{ dag.dag_id }}-{{ ts_nodash }}",
"tasks": [{"task_key": "noop", "notebook_task":
{"notebook_path": NOTEBOOK_PATH}}],
},
)
@task(task_id="verify_create_job")
def verify_create_job(job_id: int) -> int:
job = _hook()._do_api_call(("GET", "2.2/jobs/get"), {"job_id":
job_id})
params = job["settings"].get("parameters", [])
assert {"name": "env", "default": "prod"} in params, params
assert {"name": "batch_size", "default": "100"} in params, params
return job_id
run_now = DatabricksRunNowOperator(
task_id="run_now",
job_id="{{ ti.xcom_pull(task_ids='verify_create_job') }}",
params={"env": "staging", "batch_size": "42"},
wait_for_termination=False,
)
@task(task_id="verify_run_now")
def verify_run_now(databricks_run_id: int) -> None:
run = _hook()._do_api_call(("GET", "2.2/jobs/runs/get"), {"run_id":
databricks_run_id})
found = {p["name"]: p.get("value", p.get("default")) for p in
run.get("job_parameters", [])}
assert found.get("env") == "staging", found
assert found.get("batch_size") == "42", found
_hook().cancel_run(databricks_run_id)
submit_run = DatabricksSubmitRunOperator(
task_id="submit_run",
notebook_task={"notebook_path": NOTEBOOK_PATH},
new_cluster={
"spark_version": "15.4.x-scala2.12",
"node_type_id": "i3.xlarge",
"num_workers": 0,
"spark_conf": {"spark.master": "local[*]"},
"custom_tags": {"ResourceClass": "SingleNode"},
},
params={"env": "dev", "shard": "1"},
wait_for_termination=False,
)
@task(task_id="verify_submit_run")
def verify_submit_run(databricks_run_id: int) -> None:
run = _hook()._do_api_call(("GET", "2.2/jobs/runs/get"), {"run_id":
databricks_run_id})
base_params = run["tasks"][0]["notebook_task"]["base_parameters"]
# Operator-level params override DAG-level params for shared keys;
DAG-level
# keys not overridden are still inherited (here: batch_size from the
DAG).
assert base_params == {"env": "dev", "batch_size": "100", "shard":
"1"}, base_params
_hook().cancel_run(databricks_run_id)
@task(task_id="cleanup_job", trigger_rule="all_done")
def cleanup_job(job_id: int) -> None:
try:
_hook()._do_api_call(("POST", "2.2/jobs/delete"), {"job_id":
job_id})
except Exception: # noqa: BLE001
pass
job_id_xcom = verify_create_job(create_job.output)
run_now_done = verify_run_now(
databricks_run_id="{{ ti.xcom_pull(task_ids='run_now', key='run_id')
}}"
)
submit_run_done = verify_submit_run(
databricks_run_id="{{ ti.xcom_pull(task_ids='submit_run',
key='run_id') }}"
)
create_job >> job_id_xcom >> run_now >> run_now_done
submit_run >> submit_run_done
[run_now_done, submit_run_done] >> cleanup_job(job_id_xcom)
```
</details>
Run with:
```bash
export
AIRFLOW_CONN_DATABRICKS_DEFAULT='{"conn_type":"databricks","host":"https://<workspace>","password":"<token>"}'
export PR66613_NOTEBOOK_PATH=/Users/<you>/airflow-pr66613-noop # any
notebook in the workspace
airflow dags test pr66613_realenv
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
