Databricks: resolve connection asynchronously in trigger to prevent A… - #72245
Open
aman-pathak1 wants to merge 2 commits into
Open
Databricks: resolve connection asynchronously in trigger to prevent A…#72245aman-pathak1 wants to merge 2 commits into
aman-pathak1 wants to merge 2 commits into
Conversation
…syncToSync failure (apache#71525)
|
Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR addresses #71525 by making Databricks deferrable trigger execution avoid sync connection resolution inside an active asyncio event loop (the AsyncToSync failure path), by introducing/using async-safe connection lookup and async hook methods in trigger code paths.
Changes:
- Add an async-safe connection resolution method to
BaseDatabricksHookand ensure async API calls pre-resolve/cache the connection. - Replace
sync_to_async(...)wrappers in Databricks triggers with native async hook methods. - Add async hook methods for run-task pagination and cancellation operations used by triggers.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
providers/databricks/src/airflow/providers/databricks/triggers/databricks.py |
Switch trigger cancellation/status paths from sync_to_async-wrapped sync hook calls to native async hook methods. |
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py |
Add async variants for run-task retrieval and cancellation APIs so trigger code can stay fully async. |
providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py |
Add async-safe connection lookup/caching and ensure async request/token helpers pre-resolve the connection to avoid sync lookup inside the event loop. |
Suppressed comments (1)
providers/databricks/src/airflow/providers/databricks/triggers/databricks.py:289
on_kill()callsa_cancel_sql_statement(), which uses_a_do_api_call()and therefore requires an aiohttp session fromasync with self.hook. The triggerer callson_kill()after cancellingrun()(after theasync with self.hookinrun()has unwound), so_sessionmay beNonehere and cancellation can crash instead of best-effort cancelling.
Wrap the cancellation logic in async with self.hook: (and adjust unit tests that currently patch cancel_sql_statement to patch a_cancel_sql_statement).
async def on_kill(self) -> None:
"""Cancel the Databricks SQL statement when the trigger is cancelled by a user action."""
if self.statement_id:
self.log.info("Cancelling Databricks SQL statement %s.", self.statement_id)
await self.hook.a_cancel_sql_statement(self.statement_id)
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
109
to
113
| 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 |
Comment on lines
+607
to
+630
| 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 |
Author
ok |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
…syncToSync failure (#71525)
Was generative AI tooling used to co-author this PR?
{pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.