This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6da84b1a458 Add performance_target to Databricks RunNow and CreateJobs
operators (#72148)
6da84b1a458 is described below
commit 6da84b1a458916040224421b6dc2f31bd4dd5ae6
Author: PoAn Yang <[email protected]>
AuthorDate: Wed Sep 9 07:13:38 2026 +0900
Add performance_target to Databricks RunNow and CreateJobs operators
(#72148)
Signed-off-by: PoAn Yang <[email protected]>
---
.../databricks/docs/operators/jobs_create.rst | 1 +
providers/databricks/docs/operators/run_now.rst | 1 +
.../providers/databricks/operators/databricks.py | 33 ++++++++++++++++++++--
.../unit/databricks/operators/test_databricks.py | 29 +++++++++++++++++++
4 files changed, 62 insertions(+), 2 deletions(-)
diff --git a/providers/databricks/docs/operators/jobs_create.rst
b/providers/databricks/docs/operators/jobs_create.rst
index 115dcc39d09..9c2a56f170b 100644
--- a/providers/databricks/docs/operators/jobs_create.rst
+++ b/providers/databricks/docs/operators/jobs_create.rst
@@ -56,6 +56,7 @@ Currently the named parameters that
``DatabricksCreateJobsOperator`` supports ar
- ``max_concurrent_runs``
- ``git_source``
- ``access_control_list``
+ - ``performance_target``
Forwarding Airflow Dag params as Databricks job parameters
diff --git a/providers/databricks/docs/operators/run_now.rst
b/providers/databricks/docs/operators/run_now.rst
index 0b53f659b8a..8681934e2b6 100644
--- a/providers/databricks/docs/operators/run_now.rst
+++ b/providers/databricks/docs/operators/run_now.rst
@@ -46,6 +46,7 @@ All other parameters are optional and described in
documentation for ``Databrick
* ``jar_params``
* ``spark_submit_params``
* ``idempotency_token``
+* ``performance_target``
* ``repair_run``
* ``cancel_previous_runs``
diff --git
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
index 6f587f39824..009c803899a 100644
---
a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
+++
b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py
@@ -452,6 +452,15 @@ class DatabricksCreateJobsOperator(BaseOperator):
.. seealso::
This will only be used on create. In order to reset ACL consider
using the Databricks
UI.
+ :param performance_target: Optional performance mode for runs of this job
on serverless compute.
+ Either ``PERFORMANCE_OPTIMIZED`` (prioritizes fast startup and
execution) or
+ ``STANDARD`` (enables cost-efficient execution of serverless
workloads). The API drops any
+ other value instead of rejecting it, so a mistyped ``STANDARD`` raises
no error and the run
+ falls back to the more expensive default, ``PERFORMANCE_OPTIMIZED``.
This field will be
+ templated.
+
+ .. seealso::
+ https://docs.databricks.com/api/workspace/jobs/create
:param databricks_conn_id: Reference to the
:ref:`Databricks connection <howto/connection:databricks>`. (templated)
:param polling_period_seconds: Controls the rate which we poll for the
result of
@@ -488,6 +497,7 @@ class DatabricksCreateJobsOperator(BaseOperator):
"max_concurrent_runs",
"git_source",
"access_control_list",
+ "performance_target",
"databricks_conn_id",
)
# Databricks brand color (blue) under white text
@@ -511,6 +521,7 @@ class DatabricksCreateJobsOperator(BaseOperator):
max_concurrent_runs: int | None = None,
git_source: dict | None = None,
access_control_list: list[dict] | None = None,
+ performance_target: str | None = None,
databricks_conn_id: str = "databricks_default",
polling_period_seconds: int = 30,
databricks_retry_limit: int = 3,
@@ -534,6 +545,7 @@ class DatabricksCreateJobsOperator(BaseOperator):
self.max_concurrent_runs = max_concurrent_runs
self.git_source = git_source
self.access_control_list = access_control_list
+ self.performance_target = performance_target
self.databricks_conn_id = databricks_conn_id
self.polling_period_seconds = polling_period_seconds
self.databricks_retry_limit = databricks_retry_limit
@@ -555,6 +567,7 @@ class DatabricksCreateJobsOperator(BaseOperator):
"max_concurrent_runs": self.max_concurrent_runs,
"git_source": self.git_source,
"access_control_list": self.access_control_list,
+ "performance_target": self.performance_target,
}
def _get_merged_json(self) -> dict[str, Any]:
@@ -706,8 +719,10 @@ class DatabricksSubmitRunOperator(ResumableJobMixin,
BaseOperator):
supported task types are retrieved.
:param performance_target: Optional performance mode for the run on
serverless compute.
Either ``PERFORMANCE_OPTIMIZED`` (prioritizes fast startup and
execution) or
- ``STANDARD`` (enables cost-efficient execution of serverless
workloads). This field
- will be templated.
+ ``STANDARD`` (enables cost-efficient execution of serverless
workloads). The API drops any
+ other value instead of rejecting it, so a mistyped ``STANDARD`` raises
no error and the run
+ falls back to the more expensive default, ``PERFORMANCE_OPTIMIZED``.
This field will be
+ templated.
.. seealso::
https://docs.databricks.com/api/workspace/jobs/submit
@@ -1087,6 +1102,7 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
- ``jar_params``
- ``spark_submit_params``
- ``idempotency_token``
+ - ``performance_target``
- ``repair_run``
- ``databricks_repair_reason_new_settings``
- ``cancel_previous_runs``
@@ -1186,6 +1202,15 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
:param idempotency_token: an optional token that can be used to guarantee
the idempotency of job run
requests. If a run with the provided token already exists, the request
does not create a new run but
returns the ID of the existing run instead. This token must have at
most 64 characters.
+ :param performance_target: Optional performance mode for this run on
serverless compute, overriding
+ the performance target defined at the job level. Either
``PERFORMANCE_OPTIMIZED`` (prioritizes
+ fast startup and execution) or ``STANDARD`` (enables cost-efficient
execution of serverless
+ workloads). The API drops any other value instead of rejecting it, so
a mistyped ``STANDARD``
+ raises no error and the run falls back to the more expensive default,
``PERFORMANCE_OPTIMIZED``.
+ This field will be templated.
+
+ .. seealso::
+ https://docs.databricks.com/api/workspace/jobs/runnow
:param databricks_conn_id: Reference to the :ref:`Databricks connection
<howto/connection:databricks>`.
By default and in the common case this will be ``databricks_default``.
To use
token based authentication, provide the key ``token`` in the extra
field for the
@@ -1243,6 +1268,7 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
"jar_params",
"spark_submit_params",
"idempotency_token",
+ "performance_target",
"databricks_conn_id",
)
template_ext: Sequence[str] = (".json-tpl",)
@@ -1265,6 +1291,7 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
spark_submit_params: list[str] | None = None,
python_named_params: dict[str, str] | None = None,
idempotency_token: str | None = None,
+ performance_target: str | None = None,
databricks_conn_id: str = "databricks_default",
polling_period_seconds: int = 30,
databricks_retry_limit: int = 3,
@@ -1297,6 +1324,7 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
self.jar_params = jar_params
self.spark_submit_params = spark_submit_params
self.idempotency_token = idempotency_token
+ self.performance_target = performance_target
self.databricks_conn_id = databricks_conn_id
self.polling_period_seconds = polling_period_seconds
self.databricks_retry_limit = databricks_retry_limit
@@ -1325,6 +1353,7 @@ class DatabricksRunNowOperator(ResumableJobMixin,
BaseOperator):
"jar_params": self.jar_params,
"spark_submit_params": self.spark_submit_params,
"idempotency_token": self.idempotency_token,
+ "performance_target": self.performance_target,
}
def _get_merged_json(self) -> dict[str, Any]:
diff --git
a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
index 4a5dc4bd43f..990e15f1098 100644
--- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py
@@ -355,6 +355,26 @@ class TestDatabricksCreateJobsOperator:
assert expected == utils.normalise_json_content(op._get_merged_json())
+ def test_init_with_performance_target_named_parameter(self):
+ """
+ Test the initializer merges ``performance_target`` into the create
payload.
+ """
+ op = DatabricksCreateJobsOperator(
+ task_id=TASK_ID,
+ name=JOB_NAME,
+ tasks=TASKS,
+ performance_target="PERFORMANCE_OPTIMIZED",
+ )
+ expected = utils.normalise_json_content(
+ {
+ "name": JOB_NAME,
+ "tasks": TASKS,
+ "performance_target": "PERFORMANCE_OPTIMIZED",
+ }
+ )
+
+ assert expected == utils.normalise_json_content(op._get_merged_json())
+
def test_init_with_json(self):
"""
Test the initializer with json data.
@@ -1930,6 +1950,15 @@ class TestDatabricksRunNowOperator:
assert expected == utils.normalise_json_content(op._get_merged_json())
+ def test_init_with_performance_target_named_parameter(self):
+ """
+ Test the initializer merges ``performance_target`` into the run-now
payload.
+ """
+ op = DatabricksRunNowOperator(job_id=JOB_ID, task_id=TASK_ID,
performance_target="STANDARD")
+ expected = utils.normalise_json_content({"job_id": 42,
"performance_target": "STANDARD"})
+
+ assert expected == utils.normalise_json_content(op._get_merged_json())
+
def test_init_with_json(self):
"""
Test the initializer with json data.
