From 0f0fb0f3cfd7ab2470d01fd398649ea352ed9bd1 Mon Sep 17 00:00:00 2001 From: subham611 Date: Fri, 12 Apr 2024 17:09:14 +0530 Subject: [PATCH 1/4] Adds job_id as path param in update permission --- airflow/providers/databricks/hooks/databricks.py | 6 +++--- airflow/providers/databricks/operators/databricks.py | 2 +- tests/providers/databricks/operators/test_databricks.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) 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): From bd1835b6de2a48b2eed4b48e27935f8ef3baa984 Mon Sep 17 00:00:00 2001 From: subham611 Date: Fri, 12 Apr 2024 17:28:50 +0530 Subject: [PATCH 2/4] Adds try catch --- airflow/providers/databricks/hooks/databricks.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index 1a0ab8e8c6ba8..4425359fa03b0 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -663,7 +663,12 @@ def update_job_permission(self, job_id: int, json: dict[str, Any]) -> dict: :param json: payload :return: json containing permission specification """ - return self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json) + response = {} + try: + response = self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json) + except Exception as e: + print("Could not update permission ", e) + return response def test_connection(self) -> tuple[bool, str]: """Test the Databricks connectivity from UI.""" From 0de1d5875b6ebb68cdc598997c182a6d4df42e2b Mon Sep 17 00:00:00 2001 From: subham611 Date: Sat, 13 Apr 2024 14:26:05 +0530 Subject: [PATCH 3/4] print exception --- airflow/providers/databricks/hooks/databricks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index 4425359fa03b0..18319f03053a6 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -667,7 +667,7 @@ def update_job_permission(self, job_id: int, json: dict[str, Any]) -> dict: try: response = self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json) except Exception as e: - print("Could not update permission ", e) + self.log.error("Could not update permission %s", e) return response def test_connection(self) -> tuple[bool, str]: From 331dff85397edbbcaf39f90924d74fe41d3d8f70 Mon Sep 17 00:00:00 2001 From: subham611 Date: Sun, 14 Apr 2024 10:34:21 +0530 Subject: [PATCH 4/4] Removes try catch --- airflow/providers/databricks/hooks/databricks.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index 18319f03053a6..1a0ab8e8c6ba8 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -663,12 +663,7 @@ def update_job_permission(self, job_id: int, json: dict[str, Any]) -> dict: :param json: payload :return: json containing permission specification """ - response = {} - try: - response = self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json) - except Exception as e: - self.log.error("Could not update permission %s", e) - return response + 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."""