From 7847767da153a712cc977fbc425543c13816e28f Mon Sep 17 00:00:00 2001 From: Stephen Purcell Date: Tue, 11 Jun 2024 16:46:33 +0200 Subject: [PATCH 1/3] Databricks: stop including user names in `list_jobs` The user's name is not used on the Airflow side, and this argument saves the lookup, which makes the request faster. --- airflow/providers/databricks/hooks/databricks.py | 2 ++ tests/providers/databricks/hooks/test_databricks.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index ee5349add055e..4d38587e72e0c 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -256,6 +256,7 @@ def list_jobs( expand_tasks: bool = False, job_name: str | None = None, page_token: str | None = None, + include_user_names: bool = False, ) -> list[dict[str, Any]]: """ List the jobs in the Databricks Job Service. @@ -275,6 +276,7 @@ def list_jobs( payload: dict[str, Any] = { "limit": limit, "expand_tasks": expand_tasks, + "include_user_names": include_user_names, } payload["page_token"] = page_token if job_name: diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py index 0f1d2c242e456..0e478cc632d16 100644 --- a/tests/providers/databricks/hooks/test_databricks.py +++ b/tests/providers/databricks/hooks/test_databricks.py @@ -912,7 +912,7 @@ def test_list_jobs_success_single_page(self, mock_requests): mock_requests.get.assert_called_once_with( list_jobs_endpoint(HOST), json=None, - params={"limit": 25, "page_token": "", "expand_tasks": False}, + params={"limit": 25, "page_token": "", "expand_tasks": False, "include_user_names": False}, auth=HTTPBasicAuth(LOGIN, PASSWORD), headers=self.hook.user_agent_header, timeout=self.hook.timeout_seconds, From c16af02007fecf129c9354b3b7aee4a85e4ce1f8 Mon Sep 17 00:00:00 2001 From: Stephen Purcell Date: Thu, 13 Jun 2024 12:51:13 +0200 Subject: [PATCH 2/3] Update tests --- .../databricks/hooks/test_databricks.py | 32 ++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py index 0e478cc632d16..3d3e5aac4bf48 100644 --- a/tests/providers/databricks/hooks/test_databricks.py +++ b/tests/providers/databricks/hooks/test_databricks.py @@ -936,7 +936,12 @@ def test_list_jobs_success_multiple_pages(self, mock_requests): first_call_args = mock_requests.method_calls[0] assert first_call_args[1][0] == list_jobs_endpoint(HOST) - assert first_call_args[2]["params"] == {"limit": 25, "page_token": "", "expand_tasks": False} + assert first_call_args[2]["params"] == { + "limit": 25, + "page_token": "", + "expand_tasks": False, + "include_user_names": False + } second_call_args = mock_requests.method_calls[1] assert second_call_args[1][0] == list_jobs_endpoint(HOST) @@ -944,6 +949,7 @@ def test_list_jobs_success_multiple_pages(self, mock_requests): "limit": 25, "page_token": "PAGETOKEN", "expand_tasks": False, + "include_user_names": False } assert len(jobs) == 2 @@ -959,7 +965,13 @@ def test_get_job_id_by_name_success(self, mock_requests): mock_requests.get.assert_called_once_with( list_jobs_endpoint(HOST), json=None, - params={"limit": 25, "page_token": "", "expand_tasks": False, "name": JOB_NAME}, + params={ + "limit": 25, + "page_token": "", + "expand_tasks": False, + "include_user_names": False, + "name": JOB_NAME + }, auth=HTTPBasicAuth(LOGIN, PASSWORD), headers=self.hook.user_agent_header, timeout=self.hook.timeout_seconds, @@ -978,7 +990,13 @@ def test_get_job_id_by_name_not_found(self, mock_requests): mock_requests.get.assert_called_once_with( list_jobs_endpoint(HOST), json=None, - params={"limit": 25, "page_token": "", "expand_tasks": False, "name": job_name}, + params={ + "limit": 25, + "page_token": "", + "expand_tasks": False, + "include_user_names": False, + "name": job_name + }, auth=HTTPBasicAuth(LOGIN, PASSWORD), headers=self.hook.user_agent_header, timeout=self.hook.timeout_seconds, @@ -1001,7 +1019,13 @@ def test_get_job_id_by_name_raise_exception_with_duplicates(self, mock_requests) mock_requests.get.assert_called_once_with( list_jobs_endpoint(HOST), json=None, - params={"limit": 25, "page_token": "", "expand_tasks": False, "name": JOB_NAME}, + params={ + "limit": 25, + "page_token": "", + "expand_tasks": False, + "include_user_names": False, + "name": JOB_NAME + }, auth=HTTPBasicAuth(LOGIN, PASSWORD), headers=self.hook.user_agent_header, timeout=self.hook.timeout_seconds, From d21acacdaefa6bc5a37a8cb016896d721d053a11 Mon Sep 17 00:00:00 2001 From: Stephen Purcell Date: Fri, 14 Jun 2024 14:33:47 +0200 Subject: [PATCH 3/3] Run pre-commit hooks --- tests/providers/databricks/hooks/test_databricks.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py index 3d3e5aac4bf48..7e61036973182 100644 --- a/tests/providers/databricks/hooks/test_databricks.py +++ b/tests/providers/databricks/hooks/test_databricks.py @@ -940,7 +940,7 @@ def test_list_jobs_success_multiple_pages(self, mock_requests): "limit": 25, "page_token": "", "expand_tasks": False, - "include_user_names": False + "include_user_names": False, } second_call_args = mock_requests.method_calls[1] @@ -949,7 +949,7 @@ def test_list_jobs_success_multiple_pages(self, mock_requests): "limit": 25, "page_token": "PAGETOKEN", "expand_tasks": False, - "include_user_names": False + "include_user_names": False, } assert len(jobs) == 2 @@ -970,7 +970,7 @@ def test_get_job_id_by_name_success(self, mock_requests): "page_token": "", "expand_tasks": False, "include_user_names": False, - "name": JOB_NAME + "name": JOB_NAME, }, auth=HTTPBasicAuth(LOGIN, PASSWORD), headers=self.hook.user_agent_header, @@ -995,7 +995,7 @@ def test_get_job_id_by_name_not_found(self, mock_requests): "page_token": "", "expand_tasks": False, "include_user_names": False, - "name": job_name + "name": job_name, }, auth=HTTPBasicAuth(LOGIN, PASSWORD), headers=self.hook.user_agent_header, @@ -1024,7 +1024,7 @@ def test_get_job_id_by_name_raise_exception_with_duplicates(self, mock_requests) "page_token": "", "expand_tasks": False, "include_user_names": False, - "name": JOB_NAME + "name": JOB_NAME, }, auth=HTTPBasicAuth(LOGIN, PASSWORD), headers=self.hook.user_agent_header,