Copilot commented on code in PR #72245:
URL: https://github.com/apache/airflow/pull/72245#discussion_r3885825667
##########
providers/databricks/src/airflow/providers/databricks/triggers/databricks.py:
##########
@@ -108,21 +108,19 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
async def on_kill(self) -> None:
"""Cancel the Databricks run when the trigger is cancelled by a user
action."""
- from asgiref.sync import sync_to_async
-
run_id = self.run_id
if self.workflow_run_id is not None and self.databricks_task_key is
not None:
# self.run_id may be an earlier, now-terminal attempt; cancel the
task's latest attempt
Review Comment:
`on_kill()` now calls async hook methods (`a_get_run_tasks`/`a_cancel_run`)
that rely on `BaseDatabricksHook.__aenter__` having created `self._session`.
However, the triggerer invokes `on_kill()` only after cancelling
`trigger.run()`; by then the `async with self.hook` in `run()` has already
exited and `__aexit__` closes the session, so `_session` is `None` and
`_a_do_api_call()` will fail.
Wrap the `on_kill()` body in `async with self.hook:` so it has a fresh
aiohttp session during cancellation (and update unit tests that currently patch
`cancel_run/get_run_tasks` to patch the async variants).
This issue also appears on line 285 of the same file.
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -611,6 +604,31 @@ def get_run_tasks(self, run_id: int) -> list[dict[str,
Any]]:
return all_tasks
+ async def a_get_run_tasks(self, run_id: int) -> list[dict[str, Any]]:
+ """
+ Retrieve list of tasks performed by the run (async version).
+
+ :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 = await self._a_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
Review Comment:
New async hook methods (`a_get_run_tasks`, `a_cancel_run`,
`a_cancel_sql_statement`) are introduced and are now used by triggers. There
are existing unit tests for the sync counterparts, but no direct tests for
these async variants (e.g. pagination via `next_page_token`, and that the right
endpoint is passed to `_a_do_api_call`). Adding unit tests for these async
methods would prevent regressions in deferrable/triggerer paths.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
