Add performance_target to Databricks RunNow and CreateJobs operators

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

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``
 

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 <howto/connection:databricks>`. (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 <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 +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]:

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.