Uh oh!
There was an error while loading. Please reload this page.
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
a7c612b to
a878441Compare0bad94d to
fbb14acCompare
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
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
50362e6 to
687272aCompare
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 |
8e02347 to
bdd326cCompareA 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.
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.
bdd326c to
9c78611Compare
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