Update ACL during job reset

airflow/providers/databricks/hooks/databricks.py

@@ -51,6 +51,7 @@
 REPAIR_RUN_ENDPOINT = ("POST", "api/2.1/jobs/runs/repair")
 OUTPUT_RUNS_JOB_ENDPOINT = ("GET", "api/2.1/jobs/runs/get-output")
 CANCEL_ALL_RUNS_ENDPOINT = ("POST", "api/2.1/jobs/runs/cancel-all")
+UPDATE_PERMISSION_ENDPOINT = ("PATCH", "/api/2.0/permissions/jobs")
 
 INSTALL_LIBS_ENDPOINT = ("POST", "api/2.0/libraries/install")
 UNINSTALL_LIBS_ENDPOINT = ("POST", "api/2.0/libraries/uninstall")
@@ -655,6 +656,15 @@ def get_repo_by_path(self, path: str) -> str | None:
 
         return None
 
+    def update_job_permission(self, json: dict[str, Any]) -> dict:
+        """
+        Update databricks job permission.
+
+        :param json: payload
+        :return: json containing permission specification
+        """
+        return self._do_api_call(UPDATE_PERMISSION_ENDPOINT, json)
+
     def test_connection(self) -> tuple[bool, str]:
         """Test the Databricks connectivity from UI."""
         hook = DatabricksHook(databricks_conn_id=self.databricks_conn_id)

airflow/providers/databricks/operators/databricks.py

@@ -316,6 +316,10 @@ def execute(self, context: Context) -> int:
         if job_id is None:
             return self._hook.create_job(self.json)
         self._hook.reset_job(str(job_id), self.json)
+        if (access_control_list := self.json.get("access_control_list")) is not None:
+            acl_json = {"access_control_list": access_control_list}
+            self._hook.update_job_permission(normalise_json_content(acl_json))
+
         return job_id
 
 

tests/providers/databricks/operators/test_databricks.py

@@ -504,6 +504,75 @@ def test_exec_reset(self, db_mock_class):
         db_mock.reset_job.assert_called_once_with(JOB_ID, expected)
         assert JOB_ID == return_result
 
+    @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+    def test_exec_update_job_permission(self, db_mock_class):
+        """
+        Test job permission update.
+        """
+        json = {
+            "name": JOB_NAME,
+            "tags": TAGS,
+            "tasks": TASKS,
+            "job_clusters": JOB_CLUSTERS,
+            "email_notifications": EMAIL_NOTIFICATIONS,
+            "webhook_notifications": WEBHOOK_NOTIFICATIONS,
+            "timeout_seconds": TIMEOUT_SECONDS,
+            "schedule": SCHEDULE,
+            "max_concurrent_runs": MAX_CONCURRENT_RUNS,
+            "git_source": GIT_SOURCE,
+            "access_control_list": ACCESS_CONTROL_LIST,
+        }
+        op = DatabricksCreateJobsOperator(task_id=TASK_ID, json=json)
+        db_mock = db_mock_class.return_value
+        db_mock.find_job_id_by_name.return_value = JOB_ID
+
+        op.execute({})
+
+        expected = utils.normalise_json_content({"access_control_list": ACCESS_CONTROL_LIST})
+
+        db_mock_class.assert_called_once_with(
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
+            caller="DatabricksCreateJobsOperator",
+        )
+
+        db_mock.update_job_permission.assert_called_once_with(expected)
+
+    @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
+    def test_exec_update_job_permission_with_empty_acl(self, db_mock_class):
+        """
+        Test job permission update.
+        """
+        json = {
+            "name": JOB_NAME,
+            "tags": TAGS,
+            "tasks": TASKS,
+            "job_clusters": JOB_CLUSTERS,
+            "email_notifications": EMAIL_NOTIFICATIONS,
+            "webhook_notifications": WEBHOOK_NOTIFICATIONS,
+            "timeout_seconds": TIMEOUT_SECONDS,
+            "schedule": SCHEDULE,
+            "max_concurrent_runs": MAX_CONCURRENT_RUNS,
+            "git_source": GIT_SOURCE,
+        }
+        op = DatabricksCreateJobsOperator(task_id=TASK_ID, json=json)
+        db_mock = db_mock_class.return_value
+        db_mock.find_job_id_by_name.return_value = JOB_ID
+
+        op.execute({})
+
+        db_mock_class.assert_called_once_with(
+            DEFAULT_CONN_ID,
+            retry_limit=op.databricks_retry_limit,
+            retry_delay=op.databricks_retry_delay,
+            retry_args=None,
+            caller="DatabricksCreateJobsOperator",
+        )
+
+        db_mock.update_job_permission.assert_not_called()
+
 
 class TestDatabricksSubmitRunOperator:
     def test_init_with_notebook_task_named_parameters(self):