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..7e61036973182 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, @@ -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,