From 7847767da153a712cc977fbc425543c13816e28f Mon Sep 17 00:00:00 2001
From: Stephen Purcell <stephen.purcell@databricks.com>
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 <stephen.purcell@databricks.com>
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 <stephen.purcell@databricks.com>
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,
