From 2f02c4f3f3bb0932adc2784c857129dcbf0793d0 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Fri, 28 Aug 2026 16:43:51 +0900 Subject: [PATCH] Add performance_target to Databricks RunNow and CreateJobs operators Signed-off-by: PoAn Yang --- .../databricks/docs/operators/jobs_create.rst | 1 + .../databricks/docs/operators/run_now.rst | 1 + .../databricks/operators/databricks.py | 23 +++++++++++++++ .../databricks/operators/test_databricks.py | 29 +++++++++++++++++++ 4 files changed, 54 insertions(+) diff --git a/providers/databricks/docs/operators/jobs_create.rst b/providers/databricks/docs/operators/jobs_create.rst index 115dcc39d09fc..9c2a56f170b42 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 0b53f659b8afa..8681934e2b638 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 6f587f398245e..c3383a4a23ea9 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py @@ -452,6 +452,13 @@ 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). This field + will be templated. + + .. seealso:: + https://docs.databricks.com/api/workspace/jobs/create :param databricks_conn_id: Reference to the :ref:`Databricks connection `. (templated) :param polling_period_seconds: Controls the rate which we poll for the result of @@ -488,6 +495,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 +519,7 @@ def __init__( 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 +543,7 @@ def __init__( 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 +565,7 @@ def _get_named_json_parameters(self) -> dict[str, Any | None]: "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]: @@ -1087,6 +1098,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 +1198,13 @@ 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). This field will be templated. + + .. seealso:: + https://docs.databricks.com/api/workspace/jobs/runnow :param databricks_conn_id: Reference to the :ref:`Databricks connection `. 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 +1262,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 +1285,7 @@ def __init__( 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 +1318,7 @@ def __init__( 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 +1347,7 @@ def _get_named_json_parameters(self) -> dict[str, Any | None]: "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 4a5dc4bd43fc7..990e15f109840 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 @@ def test_init_with_named_parameters(self): 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 @@ def test_init_with_named_parameters(self): 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.