diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index c86391f4cdad2..1a0ab8e8c6ba8 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -51,7 +51,6 @@ 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") @@ -656,14 +655,15 @@ def get_repo_by_path(self, path: str) -> str | None: return None - def update_job_permission(self, json: dict[str, Any]) -> dict: + def update_job_permission(self, job_id: int, json: dict[str, Any]) -> dict: """ Update databricks job permission. + :param job_id: job id :param json: payload :return: json containing permission specification """ - return self._do_api_call(UPDATE_PERMISSION_ENDPOINT, json) + return self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json) def test_connection(self) -> tuple[bool, str]: """Test the Databricks connectivity from UI.""" diff --git a/airflow/providers/databricks/operators/databricks.py b/airflow/providers/databricks/operators/databricks.py index 3d95c61f6abb3..1d0d920ecca1b 100644 --- a/airflow/providers/databricks/operators/databricks.py +++ b/airflow/providers/databricks/operators/databricks.py @@ -318,7 +318,7 @@ def execute(self, context: Context) -> int: 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)) + self._hook.update_job_permission(job_id, normalise_json_content(acl_json)) return job_id diff --git a/tests/providers/databricks/operators/test_databricks.py b/tests/providers/databricks/operators/test_databricks.py index 278f95bf016b4..46e14a917ab4e 100644 --- a/tests/providers/databricks/operators/test_databricks.py +++ b/tests/providers/databricks/operators/test_databricks.py @@ -538,7 +538,7 @@ def test_exec_update_job_permission(self, db_mock_class): caller="DatabricksCreateJobsOperator", ) - db_mock.update_job_permission.assert_called_once_with(expected) + db_mock.update_job_permission.assert_called_once_with(JOB_ID, expected) @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") def test_exec_update_job_permission_with_empty_acl(self, db_mock_class):