This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0a13635c64e Cancel Snowflake queries when a user kills the deferred
task (#69635)
0a13635c64e is described below
commit 0a13635c64ef0e31f60a0fe908f4959aef6e0685
Author: Steve Ahn <[email protected]>
AuthorDate: Sun Aug 30 09:21:52 2026 -0700
Cancel Snowflake queries when a user kills the deferred task (#69635)
Killing a deferred task previously left the Snowflake statements running, so
the warehouse kept burning credits with no task left to observe it. The
worker
path already cancelled on kill; only the deferred path did not, which made
the
behaviour depend on whether the operator happened to be deferred.
Cancellation is per statement and best-effort: a failure cancelling one id
no
longer aborts the rest, since the ids most likely to error are the completed
early statements while the still-running later ones are the ones that
matter.
---
providers/snowflake/README.rst | 2 +
providers/snowflake/docs/index.rst | 2 +
providers/snowflake/pyproject.toml | 2 +
.../providers/snowflake/operators/snowflake.py | 8 +++
.../snowflake/triggers/snowflake_trigger.py | 41 +++++++++++++
.../unit/snowflake/operators/test_snowflake.py | 17 ++++++
.../unit/snowflake/triggers/test_snowflake.py | 71 ++++++++++++++++++++++
uv.lock | 3 +
8 files changed, 146 insertions(+)
diff --git a/providers/snowflake/README.rst b/providers/snowflake/README.rst
index 2a92da1081f..ddaefa33ed8 100644
--- a/providers/snowflake/README.rst
+++ b/providers/snowflake/README.rst
@@ -56,6 +56,8 @@ PIP package Version required
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.12.0``
``apache-airflow-providers-common-sql`` ``>=1.32.0``
+``asgiref`` ``>=2.3.0; python_version <
"3.14"``
+``asgiref`` ``>=3.11.1; python_version >=
"3.14"``
``pandas`` ``>=2.1.2; python_version <
"3.13"``
``pandas`` ``>=2.2.3; python_version >=
"3.13" and python_version < "3.14"``
``pandas`` ``>=2.3.3; python_version >=
"3.14"``
diff --git a/providers/snowflake/docs/index.rst
b/providers/snowflake/docs/index.rst
index d354e6ea57d..93330dc1d45 100644
--- a/providers/snowflake/docs/index.rst
+++ b/providers/snowflake/docs/index.rst
@@ -105,6 +105,8 @@ PIP package Version required
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.12.0``
``apache-airflow-providers-common-sql`` ``>=1.32.0``
+``asgiref`` ``>=2.3.0; python_version <
"3.14"``
+``asgiref`` ``>=3.11.1; python_version >=
"3.14"``
``pandas`` ``>=2.1.2; python_version <
"3.13"``
``pandas`` ``>=2.2.3; python_version >=
"3.13" and python_version < "3.14"``
``pandas`` ``>=2.3.3; python_version >=
"3.14"``
diff --git a/providers/snowflake/pyproject.toml
b/providers/snowflake/pyproject.toml
index a5a2b3afd73..8896abe6986 100644
--- a/providers/snowflake/pyproject.toml
+++ b/providers/snowflake/pyproject.toml
@@ -62,6 +62,8 @@ dependencies = [
"apache-airflow>=2.11.0",
"apache-airflow-providers-common-compat>=1.12.0",
"apache-airflow-providers-common-sql>=1.32.0",
+ "asgiref>=2.3.0; python_version < '3.14'",
+ "asgiref>=3.11.1; python_version >= '3.14'",
'pandas>=2.1.2; python_version <"3.13"',
'pandas>=2.2.3; python_version >="3.13" and python_version <"3.14"',
'pandas>=2.3.3; python_version >="3.14"',
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
index fa8b6348ca7..f716a88ef71 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
@@ -395,6 +395,9 @@ class SnowflakeSqlApiOperator(ResumableJobMixin,
SQLExecuteQueryOperator):
To set the timeout to the maximum value (604800 seconds), set
timeout to 0.
:param deferrable: Run operator in the deferrable mode.
:param snowflake_api_retry_args: An optional dictionary with arguments
passed to ``tenacity.Retrying`` & ``tenacity.AsyncRetrying`` classes.
+ :param cancel_on_kill: If True (default), cancel the running Snowflake
queries when the task is
+ killed. This applies both while the operator is running and, for a
deferred task, while it
+ waits in the triggerer.
:param durable: When ``True`` (the default), the submitted statement
handles are persisted to
task state before polling begins. A worker crash on retry reconnects
to the existing
statements instead of resubmitting the SQL. Set to ``False`` to always
submit fresh on
@@ -429,6 +432,7 @@ class SnowflakeSqlApiOperator(ResumableJobMixin,
SQLExecuteQueryOperator):
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
snowflake_api_retry_args: dict[str, Any] | None = None,
durable: bool | None = None,
+ cancel_on_kill: bool = True,
**kwargs: Any,
) -> None:
# Named here (not left to **kwargs) so default_args reaches it on every
@@ -445,6 +449,7 @@ class SnowflakeSqlApiOperator(ResumableJobMixin,
SQLExecuteQueryOperator):
self.execute_async = False
self.snowflake_api_retry_args = snowflake_api_retry_args or {}
self.deferrable = deferrable
+ self.cancel_on_kill = cancel_on_kill
self.query_ids: list[str] = []
if any([warehouse, database, role, schema, authenticator,
session_parameters]): # pragma: no cover
hook_params = kwargs.pop("hook_params", {}) # pragma: no cover
@@ -511,6 +516,7 @@ class SnowflakeSqlApiOperator(ResumableJobMixin,
SQLExecuteQueryOperator):
snowflake_conn_id=self.snowflake_conn_id,
token_life_time=self.token_life_time,
token_renewal_delta=self.token_renewal_delta,
+ cancel_on_kill=self.cancel_on_kill,
),
method_name="execute_complete",
)
@@ -637,6 +643,8 @@ class SnowflakeSqlApiOperator(ResumableJobMixin,
SQLExecuteQueryOperator):
def on_kill(self) -> None:
"""Cancel the running query."""
+ if not self.cancel_on_kill:
+ return
if self.query_ids:
self.log.info("Cancelling the query ids %s", self.query_ids)
self._hook.cancel_queries(self.query_ids)
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py
b/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py
index 233cf46e732..268fb7ac085 100644
---
a/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py
+++
b/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py
@@ -20,6 +20,8 @@ import asyncio
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Any
+from asgiref.sync import sync_to_async
+
from airflow.providers.snowflake.hooks.snowflake_sql_api import
SnowflakeSqlApiHook
from airflow.triggers.base import BaseTrigger, TriggerEvent
@@ -36,6 +38,9 @@ class SnowflakeSqlApiTrigger(BaseTrigger):
:param snowflake_conn_id: Reference to Snowflake connection id
:param token_life_time: lifetime of the JWT Token in timedelta
:param token_renewal_delta: Renewal time of the JWT Token in timedelta
+ :param cancel_on_kill: If True (default), cancel the running Snowflake
queries when the user
+ kills the deferred task (mark failed, clear, or mark success).
Requires a version of
+ ``apache-airflow`` with ``BaseTrigger.on_kill()`` support; on older
versions it is inert.
"""
def __init__(
@@ -45,6 +50,7 @@ class SnowflakeSqlApiTrigger(BaseTrigger):
snowflake_conn_id: str,
token_life_time: timedelta,
token_renewal_delta: timedelta,
+ cancel_on_kill: bool = True,
):
super().__init__()
self.poll_interval = poll_interval
@@ -52,6 +58,7 @@ class SnowflakeSqlApiTrigger(BaseTrigger):
self.snowflake_conn_id = snowflake_conn_id
self.token_life_time = token_life_time
self.token_renewal_delta = token_renewal_delta
+ self.cancel_on_kill = cancel_on_kill
def serialize(self) -> tuple[str, dict[str, Any]]:
"""Serialize SnowflakeSqlApiTrigger arguments and classpath."""
@@ -63,6 +70,7 @@ class SnowflakeSqlApiTrigger(BaseTrigger):
"snowflake_conn_id": self.snowflake_conn_id,
"token_life_time": self.token_life_time,
"token_renewal_delta": self.token_renewal_delta,
+ "cancel_on_kill": self.cancel_on_kill,
},
)
@@ -93,6 +101,39 @@ class SnowflakeSqlApiTrigger(BaseTrigger):
except Exception as e:
yield TriggerEvent({"status": "error", "message": str(e)})
+ async def on_kill(self) -> None:
+ """
+ Best-effort cancel of the running Snowflake queries when the user
kills the deferred task.
+
+ Cancellation issues one blocking request per query id, so a task with
many statements
+ against a slow warehouse can exceed the triggerer's ``[triggerer]
on_kill_timeout``
+ (default 30s); statements not cancelled within that budget may keep
running.
+ """
+ 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)()
+ 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,
+ )
+ for query_id in self.query_ids:
+ try:
+ hook.cancel_queries([query_id])
+ self.log.info("Snowflake query id %s cancelled.", query_id)
+ except Exception:
+ self.log.exception(
+ "Failed to cancel Snowflake query id %s; continuing with
the rest.", query_id
+ )
+
async def get_query_status(
self, query_id: str, hook: SnowflakeSqlApiHook | None = None
) -> dict[str, Any]:
diff --git
a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py
b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py
index 7cf8713ab7d..517691fed3e 100644
--- a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py
+++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py
@@ -459,6 +459,7 @@ class TestSnowflakeSqlApiOperator:
assert isinstance(exc.value.trigger, SnowflakeSqlApiTrigger), (
"Trigger is not a SnowflakeSqlApiTrigger"
)
+ assert exc.value.trigger.cancel_on_kill is True
def test_snowflake_sql_api_pushes_query_ids_to_xcom(
self,
@@ -753,6 +754,22 @@ class TestSnowflakeSqlApiOperator:
mock_cancel_queries.assert_not_called()
+
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.cancel_queries")
+ def test_snowflake_sql_api_on_kill_respects_cancel_on_kill_false(self,
mock_cancel_queries):
+ """on_kill does not cancel queries when cancel_on_kill is disabled."""
+ operator = SnowflakeSqlApiOperator(
+ task_id=TASK_ID,
+ snowflake_conn_id=CONN_ID,
+ sql=SQL_MULTIPLE_STMTS,
+ statement_count=4,
+ cancel_on_kill=False,
+ )
+ operator.query_ids = ["uuid1", "uuid2"]
+
+ operator.on_kill()
+
+ mock_cancel_queries.assert_not_called()
+
@pytest.mark.skipif(
not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
diff --git
a/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py
b/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py
index 42a3a07d224..757d5b997ec 100644
--- a/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py
+++ b/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py
@@ -55,8 +55,79 @@ class TestSnowflakeSqlApiTrigger:
"snowflake_conn_id": "test_conn",
"token_life_time": LIFETIME,
"token_renewal_delta": RENEWAL_DELTA,
+ "cancel_on_kill": True,
}
+ def test_snowflake_sql_trigger_serialization_cancel_on_kill_false(self):
+ """cancel_on_kill=False round-trips through serialization."""
+ trigger = SnowflakeSqlApiTrigger(
+ poll_interval=POLL_INTERVAL,
+ query_ids=QUERY_IDS,
+ snowflake_conn_id="test_conn",
+ token_life_time=LIFETIME,
+ token_renewal_delta=RENEWAL_DELTA,
+ cancel_on_kill=False,
+ )
+ _, kwargs = trigger.serialize()
+ assert kwargs["cancel_on_kill"] is False
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.triggers.snowflake_trigger.SnowflakeSqlApiHook")
+ async def test_on_kill_cancels_the_queries(self, mock_hook):
+ """on_kill() cancels the running queries when enabled and query_ids
are set."""
+ await self.TRIGGER.on_kill()
+ mock_hook.assert_called_once_with("test_conn", LIFETIME, RENEWAL_DELTA)
+
mock_hook.return_value.cancel_queries.assert_called_once_with(QUERY_IDS)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("cancel_on_kill", "query_ids"),
+ [
+ pytest.param(False, QUERY_IDS, id="disabled"),
+ pytest.param(True, [], id="no-query-ids"),
+ ],
+ )
+ @mock.patch(f"{MODULE}.triggers.snowflake_trigger.SnowflakeSqlApiHook")
+ async def test_on_kill_does_not_cancel(self, mock_hook, cancel_on_kill,
query_ids):
+ """on_kill() is a no-op (no hook built) when disabled or without
query_ids."""
+ trigger = SnowflakeSqlApiTrigger(
+ poll_interval=POLL_INTERVAL,
+ query_ids=query_ids,
+ snowflake_conn_id="test_conn",
+ token_life_time=LIFETIME,
+ token_renewal_delta=RENEWAL_DELTA,
+ cancel_on_kill=cancel_on_kill,
+ )
+ await trigger.on_kill()
+ mock_hook.assert_not_called()
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.triggers.snowflake_trigger.SnowflakeSqlApiHook")
+ async def test_on_kill_swallows_cancel_errors(self, mock_hook):
+ """on_kill() logs and swallows exceptions raised while cancelling."""
+ mock_hook.return_value.cancel_queries.side_effect =
Exception("Snowflake API error")
+ await self.TRIGGER.on_kill()
+
mock_hook.return_value.cancel_queries.assert_called_once_with(QUERY_IDS)
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.triggers.snowflake_trigger.SnowflakeSqlApiHook")
+ async def test_on_kill_cancels_remaining_after_one_fails(self, mock_hook):
+ """A failure cancelling one query id does not abort cancelling the
remaining ids."""
+ trigger = SnowflakeSqlApiTrigger(
+ poll_interval=POLL_INTERVAL,
+ query_ids=["q1", "q2", "q3"],
+ snowflake_conn_id="test_conn",
+ token_life_time=LIFETIME,
+ token_renewal_delta=RENEWAL_DELTA,
+ )
+ mock_hook.return_value.cancel_queries.side_effect = [RuntimeError("404
not found"), None, None]
+ await trigger.on_kill()
+ assert mock_hook.return_value.cancel_queries.call_args_list == [
+ mock.call(["q1"]),
+ mock.call(["q2"]),
+ mock.call(["q3"]),
+ ]
+
@pytest.mark.asyncio
@mock.patch(f"{MODULE}.triggers.snowflake_trigger.SnowflakeSqlApiTrigger.get_query_status")
@mock.patch(f"{MODULE}.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")
diff --git a/uv.lock b/uv.lock
index 92c39cc10db..afe82113f14 100644
--- a/uv.lock
+++ b/uv.lock
@@ -8116,6 +8116,7 @@ dependencies = [
{ name = "apache-airflow" },
{ name = "apache-airflow-providers-common-compat" },
{ name = "apache-airflow-providers-common-sql" },
+ { name = "asgiref" },
{ name = "pandas", version = "2.3.3", source = { registry =
"https://pypi.org/simple"; }, marker = "python_full_version < '3.11'" },
{ name = "pandas", version = "3.0.5", source = { registry =
"https://pypi.org/simple"; }, marker = "python_full_version >= '3.11'" },
{ name = "pyarrow" },
@@ -8156,6 +8157,8 @@ requires-dist = [
{ name = "apache-airflow-providers-common-sql", editable =
"providers/common/sql" },
{ name = "apache-airflow-providers-microsoft-azure", marker = "extra ==
'microsoft-azure'", editable = "providers/microsoft/azure" },
{ name = "apache-airflow-providers-openlineage", marker = "extra ==
'openlineage'", editable = "providers/openlineage" },
+ { name = "asgiref", marker = "python_full_version < '3.14'", specifier =
">=2.3.0" },
+ { name = "asgiref", marker = "python_full_version >= '3.14'", specifier =
">=3.11.1" },
{ name = "pandas", marker = "python_full_version < '3.13'", specifier =
">=2.1.2" },
{ name = "pandas", marker = "python_full_version == '3.13.*'", specifier =
">=2.2.3" },
{ name = "pandas", marker = "python_full_version >= '3.14'", specifier =
">=2.3.3" },
