potiuk commented on code in PR #69635:
URL: https://github.com/apache/airflow/pull/69635#discussion_r3615628659
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py:
##########
@@ -93,6 +101,27 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
except Exception as e:
yield TriggerEvent({"status": "error", "message": str(e)})
+ async def on_kill(self) -> None:
+ """Cancel the running Snowflake queries when the user kills the
deferred task."""
+ if not self.cancel_on_kill or not self.query_ids:
+ return
+ self.log.info("Cancelling Snowflake query ids %s", self.query_ids)
+ try:
+ await sync_to_async(self._cancel_queries)()
Review Comment:
Non-blocking, more of a docs note: the triggerer wraps this in
`asyncio.wait_for(trigger.on_kill(), timeout=_ON_CANCEL_TIMEOUT)` —
`[triggerer] on_kill_timeout`, default 30s.
Cancellation here is N sequential blocking POSTs, each with tenacity retries
on 429/503/504. A task with many statements against a slow warehouse could
exceed that budget, and because `sync_to_async` runs in a thread the remaining
cancels can't be interrupted — you'd get an `on_kill() timed out` warning while
the thread keeps going.
Benign in practice, but the docstring's "best effort" framing might be worth
extending to mention the timeout so operators know the knob exists.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @potiuk before posting
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py:
##########
@@ -20,6 +20,8 @@
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Any
+from asgiref.sync import sync_to_async
Review Comment:
This is the first and only `asgiref` usage in the Snowflake provider, and
`asgiref` isn't in this provider's `dependencies` in
`providers/snowflake/pyproject.toml`. It resolves today only transitively via
`apache-airflow`, but providers should declare what they import.
Amazon and Google both do exactly this where they use `sync_to_async`:
```toml
"asgiref>=2.3.0; python_version < '3.14'",
"asgiref>=3.11.1; python_version >= '3.14'",
```
Worth carrying over the `>=3.11.1` floor for 3.14 rather than a bare pin.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @potiuk before posting
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py:
##########
@@ -93,6 +101,27 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
except Exception as e:
yield TriggerEvent({"status": "error", "message": str(e)})
+ async def on_kill(self) -> None:
+ """Cancel the running Snowflake queries when the user kills the
deferred task."""
+ if not self.cancel_on_kill or not self.query_ids:
+ return
+ self.log.info("Cancelling Snowflake query ids %s", self.query_ids)
+ try:
+ await sync_to_async(self._cancel_queries)()
+ self.log.info("Snowflake query ids %s cancelled.", self.query_ids)
+ except Exception:
+ self.log.exception(
+ "Failed to cancel Snowflake query ids %s. They may still be
running.", self.query_ids
+ )
+
+ def _cancel_queries(self) -> None:
+ hook = SnowflakeSqlApiHook(
+ self.snowflake_conn_id,
+ self.token_life_time,
+ self.token_renewal_delta,
+ )
+ hook.cancel_queries(self.query_ids)
Review Comment:
`SnowflakeSqlApiHook.cancel_queries` loops `_cancel_sql_api_query_execution`
per id with no per-query error handling, and `on_kill` catches only at the
outermost level — so the first id that raises aborts cancellation of every id
after it.
That matters most in exactly the case this PR targets: `run()` polls
statements sequentially and never prunes `self.query_ids`, so on a
multi-statement operator the earlier statements have typically already
completed by kill time. If cancelling a completed statement raises, the
still-running later statements — the ones actually burning credits — never get
cancelled.
Mitigating factor: `_process_response` returns a dict for HTTP 422 rather
than raising, so if Snowflake answers 422 for "not running" the loop survives
fine. But a 404 on a reaped handle, or a 401 on token expiry, would raise.
Did your live test cover killing a *multi-statement* task where some
statements had already finished? If 422 is what Snowflake actually returns
there, this is a non-issue and worth a one-line comment saying so. If not, a
per-query `try`/`except` in `_cancel_queries` would make it genuinely
best-effort. I didn't want to assert a bug I couldn't reproduce without live
credentials.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @potiuk before posting
--
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]
