This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 069a973de5b Fix Databricks hook dropping tasks beyond the first page
of a run (#72304)
069a973de5b is described below
commit 069a973de5b3ca6be2da70ce11093282a861634a
Author: Noritaka Sekiyama <[email protected]>
AuthorDate: Thu Sep 10 05:44:32 2026 +0900
Fix Databricks hook dropping tasks beyond the first page of a run (#72304)
The Jobs API 2.2 returns at most 100 entries of a run's tasks and
job_clusters per page and hands back a token for the rest, but get_run
read only the first page. Callers that inspect the task list, such as the
failed-task error extraction behind a task's failure message, therefore
saw an arbitrary subset of it: the entries do not come back in the order
they were declared, so which ones went missing was not predictable.
---
.../providers/databricks/hooks/databricks.py | 57 ++++++++++++----------
.../tests/unit/databricks/hooks/test_databricks.py | 48 ++++++++++++++++++
2 files changed, 80 insertions(+), 25 deletions(-)
diff --git
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py
index 2223b5a5a4b..f21882a5763 100644
--- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py
+++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py
@@ -71,6 +71,19 @@ SPARK_VERSIONS_ENDPOINT = ("GET",
"2.1/clusters/spark-versions")
SQL_STATEMENTS_ENDPOINT = "2.0/sql/statements"
SQL_WAREHOUSES_ENDPOINT = "2.0/sql/warehouses"
+# ``2.2/jobs/runs/get`` returns at most 100 entries of these nested arrays per
page and hands back a
+# ``next_page_token`` for the rest, while repeating every other field on each
page. The order of the
+# entries is not the order they were declared in, so a truncated response
drops an arbitrary subset.
+PAGINATED_RUN_FIELDS = ("tasks", "job_clusters")
+
+
+def _merge_run_page(run: dict[str, Any], page: dict[str, Any]) -> None:
+ """Append the paginated arrays of ``page`` to those already collected in
``run``."""
+ for field in PAGINATED_RUN_FIELDS:
+ entries = page.get(field)
+ if entries:
+ run[field] = [*run.get(field, []), *entries]
+
class RunLifeCycleState(Enum):
"""
@@ -593,45 +606,39 @@ class DatabricksHook(BaseDatabricksHook):
:param run_id: id of the run
:return: A list of tasks
"""
- has_more = True
- all_tasks = []
- page_token = ""
- json: dict[str, Any] = {"run_id": run_id}
-
- while has_more:
- if page_token:
- json = {**json, "page_token": page_token}
- response = self._do_api_call(GET_RUN_ENDPOINT, json)
- tasks = response.get("tasks", [])
- all_tasks += tasks
- if "next_page_token" in response:
- page_token = response["next_page_token"]
- else:
- has_more = False
-
- return all_tasks
+ return self.get_run(run_id).get("tasks", [])
def get_run(self, run_id: int) -> dict[str, Any]:
"""
Retrieve run information.
:param run_id: id of the run
- :return: state of the run
+ :return: state of the run, with every page of the paginated arrays
collected
"""
- json = {"run_id": run_id}
- response = self._do_api_call(GET_RUN_ENDPOINT, json)
- return response
+ json: dict[str, Any] = {"run_id": run_id}
+ run = self._do_api_call(GET_RUN_ENDPOINT, json)
+ page_token = run.pop("next_page_token", None)
+ while page_token:
+ page = self._do_api_call(GET_RUN_ENDPOINT, {**json, "page_token":
page_token})
+ _merge_run_page(run, page)
+ page_token = page.get("next_page_token")
+ return run
async def a_get_run(self, run_id: int) -> dict[str, Any]:
"""
Async version of `get_run`.
:param run_id: id of the run
- :return: state of the run
+ :return: state of the run, with every page of the paginated arrays
collected
"""
- json = {"run_id": run_id}
- response = await self._a_do_api_call(GET_RUN_ENDPOINT, json)
- return response
+ json: dict[str, Any] = {"run_id": run_id}
+ run = await self._a_do_api_call(GET_RUN_ENDPOINT, json)
+ page_token = run.pop("next_page_token", None)
+ while page_token:
+ page = await self._a_do_api_call(GET_RUN_ENDPOINT, {**json,
"page_token": page_token})
+ _merge_run_page(run, page)
+ page_token = page.get("next_page_token")
+ return run
def get_run_state_str(self, run_id: int) -> str:
"""
diff --git
a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
index c1959cd2359..38d0bd1da10 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
@@ -750,6 +750,36 @@ class TestDatabricksHook:
assert len(tasks) == 2
assert tasks == GET_RUN_RESPONSE["tasks"] * 2
+ @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests")
+ def test_get_run_collects_every_page(self, mock_requests):
+ mock_requests.codes.ok = 200
+ mock_requests.get.side_effect = [
+ create_successful_response_mock(
+ {
+ **GET_RUN_RESPONSE,
+ "tasks": [{"task_key": "first"}],
+ "job_clusters": [{"job_cluster_key": "jc_a"}],
+ "next_page_token": "PAGETOKEN",
+ }
+ ),
+ create_successful_response_mock(
+ {
+ **GET_RUN_RESPONSE,
+ "tasks": [{"task_key": "second"}],
+ "job_clusters": [{"job_cluster_key": "jc_b"}],
+ }
+ ),
+ ]
+
+ run = self.hook.get_run(RUN_ID)
+
+ assert mock_requests.get.call_count == 2
+ assert mock_requests.method_calls[1][2]["params"] == {"run_id":
RUN_ID, "page_token": "PAGETOKEN"}
+ assert run["tasks"] == [{"task_key": "first"}, {"task_key": "second"}]
+ assert run["job_clusters"] == [{"job_cluster_key": "jc_a"},
{"job_cluster_key": "jc_b"}]
+ assert "next_page_token" not in run
+ assert run["state"] == GET_RUN_RESPONSE["state"]
+
@mock.patch("airflow.providers.databricks.hooks.databricks_base.requests")
def test_cancel_run(self, mock_requests):
mock_requests.post.return_value.json.return_value = GET_RUN_RESPONSE
@@ -2132,6 +2162,24 @@ class TestDatabricksHookAsyncMethods:
timeout=self.hook.timeout_seconds,
)
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
+ async def test_a_get_run_collects_every_page(self, mock_get):
+ mock_get.return_value.__aenter__.return_value.json = AsyncMock(
+ side_effect=[
+ {**GET_RUN_RESPONSE, "tasks": [{"task_key": "first"}],
"next_page_token": "PAGETOKEN"},
+ {**GET_RUN_RESPONSE, "tasks": [{"task_key": "second"}]},
+ ]
+ )
+
+ async with self.hook:
+ run = await self.hook.a_get_run(RUN_ID)
+
+ assert mock_get.call_count == 2
+ assert mock_get.call_args_list[1].kwargs["json"] == {"run_id": RUN_ID,
"page_token": "PAGETOKEN"}
+ assert run["tasks"] == [{"task_key": "first"}, {"task_key": "second"}]
+ assert "next_page_token" not in run
+
@pytest.mark.asyncio
@mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
async def test_get_cluster_state(self, mock_get):
