Databricks: stop including user names in list_jobs (apache#40178)

* 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.
romsharon98 · Jul 26, 2024 · 41d3338 · 41d3338
1 parent cf74264
commit 41d3338
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git 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
@@ -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,
@@ -936,14 +936,20 @@ 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)
         assert second_call_args[2]["params"] == {
             "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,