steveahnahn commented on code in PR #69635:
URL: https://github.com/apache/airflow/pull/69635#discussion_r3652688343
##########
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:
No, my live check did not cover that case.
The SQL API reference does document the cancel endpoint returning 422
(QueryFailureStatus, code 000709 "not found") rather than 404. That matches the
hook's _should_raise_for_status (>= 400 and != 422), so the loop would survive.
The token-expiry case is also mostly moot here, since _cancel_queries builds a
fresh hook and mints a fresh JWT at cancel time rather than reusing the
poll-time token.
Since that is documented rather than live-verified, I did not want to rely
on it. _cancel_queries now cancels per id inside a try/except, so any single
failure cannot stand the still-running statements after it in the list.
--
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]
