Cancel Snowflake queries when a user kills the deferred task - #69635
Cancel Snowflake queries when a user kills the deferred task#69635steveahnahn wants to merge 3 commits into
Conversation
boring-cyborg
Bot
added
area:providers
provider:snowflake
a7c612b to
a878441
Compare
potiuk
added
the
ready for maintainer review
0bad94d to
fbb14ac
Compare
potiuk
left a comment
potiuk
left a comment
There was a problem hiding this comment.
Nice change — this closes a real gap, and threading cancel_on_kill through the defer site so the trigger inherits the operator's setting is the right shape.
I verified locally: the 57 trigger + operator tests pass, the async def on_kill signature matches BaseTrigger.on_kill (the triggerer awaits it, so a sync override would have silently never run), and the positional SnowflakeSqlApiHook(conn_id, token_life_time, token_renewal_delta) args match the constructor. Building the hook inside the sync_to_async worker rather than on the event loop is the correct call. Serialization round-trip coverage and the live-verification screenshots are appreciated.
One thing to fix and two I'd like your read on — left inline. Only the asgiref one is blocking.
Drafted-by: Claude Code (Opus 4.8); reviewed by @potiuk before posting
A deferred SnowflakeSqlApiOperator parks its running query ids in the triggerer, so the operator's own on_kill no longer runs once the task is deferred. When a user marks that task failed, clears it, or marks it success, SnowflakeSqlApiTrigger had no on_kill hook, so the Snowflake statements kept executing on the warehouse, burning compute credits, even though the operator already cancels the queries on kill in the non-deferred path. This adds on_kill to SnowflakeSqlApiTrigger to cancel the running query ids when the user acts on the deferred task, matching the behaviour already shipped for the EMR, Dataproc, BigQuery, and Dataflow triggers. The Snowflake SQL API cancel is a blocking POST with no async variant, so it runs through sync_to_async off the triggerer event loop, and the hook is built inside that worker so no connection work touches the loop. A cancel_on_kill flag on both the operator and the trigger lets users opt out.
44bf75c to
50362e6
Compare
A single failing cancel (a reaped statement handle, a transient error) would otherwise stop the remaining ids from being cancelled, leaving the still-running statements that are actually consuming credits untouched. Declare asgiref, which the trigger imports directly rather than relying on it resolving transitively through apache-airflow.
50362e6 to
687272a
Compare
|
Quickest fix: git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-leaseAutomated nudge — ignore if you're not ready to rebase. This comment is updated in place on future |
2 participants
Problem
SnowflakeSqlApiOperator.on_killcancels the running queries, but afterdefer()the worker is gone and it can never run.SnowflakeSqlApiTriggerhad noon_kill(base default is a no op), so killing a deferred task leaves the statements executing on the warehouse, burning credits.Change
on_killtoSnowflakeSqlApiTrigger: cancels the running query ids.POSTwith no async variant, soon_killrunscancel_queriesthroughsync_to_async, building the hook inside that worker so no connection work touches the event loop.cancel_on_killflag (defaultTrue) on operator and trigger; operator threads it into the trigger at the defer site.Live verification
deferred, trigger polling (GET /api/v2/statements/{id}returning 202).POST /api/v2/statements/{id}/canceland the statement transitioned toABORTED. Without the change it stays running.Tests
cancel_on_kill=Falseguard, flag threaded into the trigger on defer.Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Fable 5) following the guidelines