From 7f143ff6d1c67c4d86c21c100cad3c366931f870 Mon Sep 17 00:00:00 2001 From: Steve Ahn <38049807+steveahnahn@users.noreply.github.com> Date: Wed, 8 Jul 2026 18:38:22 -0700 Subject: [PATCH 1/3] Cancel Snowflake queries when a user kills the deferred task 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. --- .../snowflake/operators/snowflake.py | 8 +++ .../snowflake/triggers/snowflake_trigger.py | 29 +++++++++++ .../snowflake/operators/test_snowflake.py | 17 ++++++ .../unit/snowflake/triggers/test_snowflake.py | 52 +++++++++++++++++++ 4 files changed, 106 insertions(+) diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py index 8e3cda63ef3d7..288ea6328c507 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py @@ -380,6 +380,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 @@ -413,6 +416,7 @@ def __init__( timeout: int | None = None, deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), snowflake_api_retry_args: dict[str, Any] | None = None, + cancel_on_kill: bool = True, **kwargs: Any, ) -> None: self.snowflake_conn_id = snowflake_conn_id @@ -425,6 +429,7 @@ def __init__( 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 @@ -491,6 +496,7 @@ def execute(self, context: Context) -> None: 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", ) @@ -617,6 +623,8 @@ def execute_complete(self, context: Context, event: dict[str, str | list[str]] | 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 e460ea7840ce4..30993a661499d 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 @@ 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 @@ def __init__( 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 @@ def __init__( 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 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "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,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) + 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 d37f420f2a66f..f37785ab7c8fe 100644 --- a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py @@ -456,6 +456,7 @@ def test_snowflake_sql_api_execute_operator_async( 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, @@ -750,6 +751,22 @@ def test_snowflake_sql_api_on_kill_no_queries(self, mock_cancel_queries): 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 42a3a07d224de..5eb941e9e0a77 100644 --- a/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py @@ -55,8 +55,60 @@ def test_snowflake_sql_trigger_serialization(self): "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.SnowflakeSqlApiTrigger.get_query_status") @mock.patch(f"{MODULE}.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async") From f3a3ce32c91af41a07b5bb385bfdbdbf1d611bf2 Mon Sep 17 00:00:00 2001 From: Steve Ahn <38049807+steveahnahn@users.noreply.github.com> Date: Thu, 16 Jul 2026 16:10:33 -0700 Subject: [PATCH 2/3] Re-trigger CI after transient CodeQL and release-tooling check failures From 8e0234719fa5eaf953eb017079a3d7f5f800d0a5 Mon Sep 17 00:00:00 2001 From: Steve Ahn <38049807+steveahnahn@users.noreply.github.com> Date: Sat, 25 Jul 2026 17:47:52 -0700 Subject: [PATCH 3/3] Make Snowflake query cancellation best-effort per statement 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. --- providers/snowflake/README.rst | 2 ++ providers/snowflake/docs/index.rst | 2 ++ providers/snowflake/pyproject.toml | 2 ++ .../snowflake/triggers/snowflake_trigger.py | 18 +++++++++++++++--- .../unit/snowflake/triggers/test_snowflake.py | 19 +++++++++++++++++++ uv.lock | 5 ++++- 6 files changed, 44 insertions(+), 4 deletions(-) diff --git a/providers/snowflake/README.rst b/providers/snowflake/README.rst index 680cb0d8cf374..a21a1dc89deef 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,<3; python_version < "3.13"`` ``pandas`` ``>=2.2.3,<3; python_version >= "3.13" and python_version < "3.14"`` ``pandas`` ``>=2.3.3,<3; python_version >= "3.14"`` diff --git a/providers/snowflake/docs/index.rst b/providers/snowflake/docs/index.rst index 1fb9aa7c2ec7a..153cddad99b90 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,<3; python_version < "3.13"`` ``pandas`` ``>=2.2.3,<3; python_version >= "3.13" and python_version < "3.14"`` ``pandas`` ``>=2.3.3,<3; python_version >= "3.14"`` diff --git a/providers/snowflake/pyproject.toml b/providers/snowflake/pyproject.toml index 5503d6e48cc3f..5fd958eb4fa48 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 3 changes the dtypes a DataFrame reads back with, so DataFrame XComs do not # round trip; capped until pandas 3 support lands. # Tracked at https://github.com/apache/airflow/pull/70558 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 30993a661499d..807e3e1f321d6 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py +++ b/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py @@ -102,13 +102,18 @@ async def run(self) -> AsyncIterator[TriggerEvent]: 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.""" + """ + 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)() - 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 @@ -120,7 +125,14 @@ def _cancel_queries(self) -> None: self.token_life_time, self.token_renewal_delta, ) - hook.cancel_queries(self.query_ids) + 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 diff --git a/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py b/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py index 5eb941e9e0a77..757d5b997ecdd 100644 --- a/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py @@ -109,6 +109,25 @@ async def test_on_kill_swallows_cancel_errors(self, mock_hook): 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 0c7c14c1c52b0..08aabd74ee6bf 100644 --- a/uv.lock +++ b/uv.lock @@ -2712,7 +2712,7 @@ requires-dist = [ { name = "types-deprecated", marker = "extra == 'mypy'", specifier = ">=1.2.9.20240311" }, { name = "types-docutils", marker = "extra == 'mypy'", specifier = ">=0.21.0.20240704" }, { name = "types-markdown", marker = "extra == 'mypy'", specifier = ">=3.6.0.20240316" }, - { name = "types-paramiko", marker = "extra == 'mypy'", specifier = ">=4.0.0.20260402,<5.0.0" }, + { name = "types-paramiko", marker = "extra == 'mypy'", specifier = ">=4.0.0.20260402" }, { name = "types-protobuf", marker = "extra == 'mypy'", specifier = ">=5.26.0.20240422" }, { name = "types-pymysql", marker = "extra == 'mypy'", specifier = ">=1.1.0.20240425" }, { name = "types-python-dateutil", marker = "extra == 'mypy'", specifier = ">=2.9.0.20240316" }, @@ -7786,6 +7786,7 @@ dependencies = [ { name = "apache-airflow" }, { name = "apache-airflow-providers-common-compat" }, { name = "apache-airflow-providers-common-sql" }, + { name = "asgiref" }, { name = "pandas" }, { name = "pyarrow" }, { name = "setuptools" }, @@ -7824,6 +7825,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,<3" }, { name = "pandas", marker = "python_full_version == '3.13.*'", specifier = ">=2.2.3,<3" }, { name = "pandas", marker = "python_full_version >= '3.14'", specifier = ">=2.3.3,<3" },