apache-airflow-providers-snowflake 6.3.0__py3-none-any.whl → 6.8.0rc1__py3-none-any.whl

airflow/providers/snowflake/hooks/snowflake_sql_api.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import base64
+import time
 import uuid
 import warnings
 from datetime import timedelta
@@ -25,10 +26,21 @@ from typing import Any
 
 import aiohttp
 import requests
+from aiohttp import ClientConnectionError, ClientResponseError
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
-
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from requests.exceptions import ConnectionError, HTTPError, Timeout
+from tenacity import (
+    AsyncRetrying,
+    Retrying,
+    before_sleep_log,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
 from airflow.providers.snowflake.utils.sql_api_generate_jwt import JWTGenerator
 
@@ -65,6 +77,7 @@ class SnowflakeSqlApiHook(SnowflakeHook):
     :param token_life_time: lifetime of the JWT Token in timedelta
     :param token_renewal_delta: Renewal time of the JWT Token in timedelta
     :param deferrable: Run operator in the deferrable mode.
+    :param api_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` & ``tenacity.AsyncRetrying`` classes.
     """
 
     LIFETIME = timedelta(minutes=59)  # The tokens will have a 59 minute lifetime
@@ -75,15 +88,27 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         snowflake_conn_id: str,
         token_life_time: timedelta = LIFETIME,
         token_renewal_delta: timedelta = RENEWAL_DELTA,
+        api_retry_args: dict[Any, Any] | None = None,  # Optional retry arguments passed to tenacity.retry
         *args: Any,
         **kwargs: Any,
     ):
         self.snowflake_conn_id = snowflake_conn_id
         self.token_life_time = token_life_time
         self.token_renewal_delta = token_renewal_delta
+
         super().__init__(snowflake_conn_id, *args, **kwargs)
         self.private_key: Any = None
 
+        self.retry_config = {
+            "retry": retry_if_exception(self._should_retry_on_error),
+            "wait": wait_exponential(multiplier=1, min=1, max=60),
+            "stop": stop_after_attempt(5),
+            "before_sleep": before_sleep_log(self.log, log_level=20),  # type: ignore[arg-type]
+            "reraise": True,
+        }
+        if api_retry_args:
+            self.retry_config.update(api_retry_args)
+
     def get_private_key(self) -> None:
         """Get the private key from snowflake connection."""
         conn = self.get_connection(self.snowflake_conn_id)
@@ -137,6 +162,7 @@ class SnowflakeSqlApiHook(SnowflakeHook):
             When executing the statement, Snowflake replaces placeholders (? and :name) in
             the statement with these specified values.
         """
+        self.query_ids = []
         conn_config = self._get_conn_params
 
         req_id = uuid.uuid4()
@@ -167,13 +193,8 @@ class SnowflakeSqlApiHook(SnowflakeHook):
                 "query_tag": query_tag,
             },
         }
-        response = requests.post(url, json=data, headers=headers, params=params)
-        try:
-            response.raise_for_status()
-        except requests.exceptions.HTTPError as e:  # pragma: no cover
-            msg = f"Response: {e.response.content.decode()} Status Code: {e.response.status_code}"
-            raise AirflowException(msg)
-        json_response = response.json()
+
+        _, json_response = self._make_api_call_with_retries("POST", url, headers, params, data)
         self.log.info("Snowflake SQL POST API response: %s", json_response)
         if "statementHandles" in json_response:
             self.query_ids = json_response["statementHandles"]
@@ -222,25 +243,37 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         }
         return headers
 
-    def get_oauth_token(self, conn_config: dict[str, Any] | None = None) -> str:
+    def get_oauth_token(
+        self,
+        conn_config: dict[str, Any] | None = None,
+        token_endpoint: str | None = None,
+        grant_type: str = "refresh_token",
+    ) -> str:
         """Generate temporary OAuth access token using refresh token in connection details."""
         warnings.warn(
             "This method is deprecated. Please use `get_oauth_token` method from `SnowflakeHook` instead. ",
             AirflowProviderDeprecationWarning,
             stacklevel=2,
         )
-        return super().get_oauth_token(conn_config=conn_config)
+        return super().get_oauth_token(
+            conn_config=conn_config, token_endpoint=token_endpoint, grant_type=grant_type
+        )
 
-    def get_request_url_header_params(self, query_id: str) -> tuple[dict[str, Any], dict[str, Any], str]:
+    def get_request_url_header_params(
+        self, query_id: str, url_suffix: str | None = None
+    ) -> tuple[dict[str, Any], dict[str, Any], str]:
         """
         Build the request header Url with account name identifier and query id from the connection params.
 
         :param query_id: statement handles query ids for the individual statements.
+        :param url_suffix: Optional path suffix to append to the URL. Must start with '/', e.g. '/cancel' or '/result'.
         """
         req_id = uuid.uuid4()
         header = self.get_headers()
         params = {"requestId": str(req_id)}
         url = f"{self.account_identifier}.snowflakecomputing.com/api/v2/statements/{query_id}"
+        if url_suffix:
+            url += url_suffix
         return header, params, url
 
     def check_query_output(self, query_ids: list[str]) -> None:
@@ -251,20 +284,31 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         """
         for query_id in query_ids:
             header, params, url = self.get_request_url_header_params(query_id)
-            try:
-                response = requests.get(url, headers=header, params=params)
-                response.raise_for_status()
-                self.log.info(response.json())
-            except requests.exceptions.HTTPError as e:
-                msg = f"Response: {e.response.content.decode()}, Status Code: {e.response.status_code}"
-                raise AirflowException(msg)
+            _, response_json = self._make_api_call_with_retries(
+                method="GET", url=url, headers=header, params=params
+            )
+            self.log.info(response_json)
 
     def _process_response(self, status_code, resp):
         self.log.info("Snowflake SQL GET statements status API response: %s", resp)
         if status_code == 202:
             return {"status": "running", "message": "Query statements are still running"}
         if status_code == 422:
-            return {"status": "error", "message": resp["message"]}
+            error_message = resp.get("message", "Unknown error occurred")
+            error_details = []
+            if code := resp.get("code"):
+                error_details.append(f"Code: {code}")
+            if sql_state := resp.get("sqlState"):
+                error_details.append(f"SQL State: {sql_state}")
+            if statement_handle := resp.get("statementHandle"):
+                error_details.append(f"Statement Handle: {statement_handle}")
+
+            if error_details:
+                enhanced_message = f"{error_message} ({', '.join(error_details)})"
+            else:
+                enhanced_message = error_message
+
+            return {"status": "error", "message": enhanced_message}
         if status_code == 200:
             if resp_statement_handles := resp.get("statementHandles"):
                 statement_handles = resp_statement_handles
@@ -287,11 +331,83 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         """
         self.log.info("Retrieving status for query id %s", query_id)
         header, params, url = self.get_request_url_header_params(query_id)
-        response = requests.get(url, params=params, headers=header)
-        status_code = response.status_code
-        resp = response.json()
+        status_code, resp = self._make_api_call_with_retries("GET", url, header, params)
         return self._process_response(status_code, resp)
 
+    def wait_for_query(
+        self, query_id: str, raise_error: bool = False, poll_interval: int = 5, timeout: int = 60
+    ) -> dict[str, str | list[str]]:
+        """
+        Wait for query to finish either successfully or with error.
+
+        :param query_id: statement handle id for the individual statement.
+        :param raise_error: whether to raise an error if the query failed.
+        :param poll_interval: time (in seconds) between checking the query status.
+        :param timeout: max time (in seconds) to wait for the query to finish before raising a TimeoutError.
+
+        :raises RuntimeError: If the query status is 'error' and `raise_error` is True.
+        :raises TimeoutError: If the query doesn't finish within the specified timeout.
+        """
+        start_time = time.time()
+
+        while True:
+            response = self.get_sql_api_query_status(query_id=query_id)
+            self.log.debug("Query status `%s`", response["status"])
+
+            if time.time() - start_time > timeout:
+                raise TimeoutError(
+                    f"Query `{query_id}` did not finish within the timeout period of {timeout} seconds."
+                )
+
+            if response["status"] != "running":
+                self.log.info("Query status `%s`", response["status"])
+                break
+
+            time.sleep(poll_interval)
+
+        if response["status"] == "error" and raise_error:
+            raise RuntimeError(response["message"])
+
+        return response
+
+    def get_result_from_successful_sql_api_query(self, query_id: str) -> list[dict[str, Any]]:
+        """
+        Based on the query id HTTP requests are made to snowflake SQL API and return result data.
+
+        :param query_id: statement handle id for the individual statement.
+
+        :raises RuntimeError: If the query status is not 'success'.
+        """
+        self.log.info("Retrieving data for query id %s", query_id)
+        header, params, url = self.get_request_url_header_params(query_id)
+        status_code, response = self._make_api_call_with_retries("GET", url, header, params)
+
+        if (query_status := self._process_response(status_code, response)["status"]) != "success":
+            msg = f"Query must have status `success` to retrieve data; got `{query_status}`."
+            raise RuntimeError(msg)
+
+        # Below fields should always be present in response, but added some safety checks
+        data = response.get("data", [])
+        if not data:
+            self.log.warning("No data found in the API response.")
+            return []
+        metadata = response.get("resultSetMetaData", {})
+        col_names = [row["name"] for row in metadata.get("rowType", [])]
+        if not col_names:
+            self.log.warning("No column metadata found in the API response.")
+            return []
+
+        num_partitions = len(metadata.get("partitionInfo", []))
+        if num_partitions > 1:
+            self.log.debug("Result data is returned as multiple partitions. Will perform additional queries.")
+            url += "?partition="
+            for partition_no in range(1, num_partitions):  # First partition was already returned
+                self.log.debug("Querying for partition no. %s", partition_no)
+                _, response = self._make_api_call_with_retries("GET", url + str(partition_no), header, params)
+                data.extend(response.get("data", []))
+
+        return [dict(zip(col_names, row)) for row in data]  # Merged column names with data
+
     async def get_sql_api_query_status_async(self, query_id: str) -> dict[str, str | list[str]]:
         """
         Based on the query id async HTTP request is made to snowflake SQL API and return response.
@@ -300,10 +416,91 @@ class SnowflakeSqlApiHook(SnowflakeHook):
         """
         self.log.info("Retrieving status for query id %s", query_id)
         header, params, url = self.get_request_url_header_params(query_id)
-        async with (
-            aiohttp.ClientSession(headers=header) as session,
-            session.get(url, params=params) as response,
+        status_code, resp = await self._make_api_call_with_retries_async("GET", url, header, params)
+        return self._process_response(status_code, resp)
+
+    def _cancel_sql_api_query_execution(self, query_id: str) -> dict[str, str | list[str]]:
+        self.log.info("Cancelling query id %s", query_id)
+        header, params, url = self.get_request_url_header_params(query_id, "/cancel")
+        status_code, resp = self._make_api_call_with_retries("POST", url, header, params)
+        return self._process_response(status_code, resp)
+
+    def cancel_queries(self, query_ids: list[str]) -> None:
+        for query_id in query_ids:
+            self._cancel_sql_api_query_execution(query_id)
+
+    @staticmethod
+    def _should_retry_on_error(exception) -> bool:
+        """
+        Determine if the exception should trigger a retry based on error type and status code.
+
+        Retries on HTTP errors 429 (Too Many Requests), 503 (Service Unavailable),
+        and 504 (Gateway Timeout) as recommended by Snowflake error handling docs.
+        Retries on connection errors and timeouts.
+
+        :param exception: The exception to check
+        :return: True if the request should be retried, False otherwise
+        """
+        if isinstance(exception, HTTPError):
+            return exception.response.status_code in [429, 503, 504]
+        if isinstance(exception, ClientResponseError):
+            return exception.status in [429, 503, 504]
+        if isinstance(
+            exception,
+            ConnectionError | Timeout | ClientConnectionError,
         ):
-            status_code = response.status
-            resp = await response.json()
-            return self._process_response(status_code, resp)
+            return True
+        return False
+
+    def _make_api_call_with_retries(
+        self, method: str, url: str, headers: dict, params: dict | None = None, json: dict | None = None
+    ):
+        """
+        Make an API call to the Snowflake SQL API with retry logic for specific HTTP errors.
+
+        Error handling implemented based on Snowflake error handling docs:
+        https://docs.snowflake.com/en/developer-guide/sql-api/handling-errors
+
+        :param method: The HTTP method to use for the API call.
+        :param url: The URL for the API endpoint.
+        :param headers: The headers to include in the API call.
+        :param params: (Optional) The query parameters to include in the API call.
+        :param json: (Optional) The data to include in the API call.
+        :return: The response object from the API call.
+        """
+        with requests.Session() as session:
+            for attempt in Retrying(**self.retry_config):  # type: ignore
+                with attempt:
+                    if method.upper() in ("GET", "POST"):
+                        response = session.request(
+                            method=method.lower(), url=url, headers=headers, params=params, json=json
+                        )
+                    else:
+                        raise ValueError(f"Unsupported HTTP method: {method}")
+                    response.raise_for_status()
+                    return response.status_code, response.json()
+
+    async def _make_api_call_with_retries_async(self, method, url, headers, params=None):
+        """
+        Make an API call to the Snowflake SQL API asynchronously with retry logic for specific HTTP errors.
+
+        Error handling implemented based on Snowflake error handling docs:
+        https://docs.snowflake.com/en/developer-guide/sql-api/handling-errors
+
+        :param method: The HTTP method to use for the API call. Only GET is supported as  is synchronous.
+        :param url: The URL for the API endpoint.
+        :param headers: The headers to include in the API call.
+        :param params: (Optional) The query parameters to include in the API call.
+        :return: The response object from the API call.
+        """
+        async with aiohttp.ClientSession(headers=headers) as session:
+            async for attempt in AsyncRetrying(**self.retry_config):
+                with attempt:
+                    if method.upper() == "GET":
+                        async with session.request(method=method.lower(), url=url, params=params) as response:
+                            response.raise_for_status()
+                            # Return status and json content for async processing
+                            content = await response.json()
+                            return response.status, content
+                    else:
+                        raise ValueError(f"Unsupported HTTP method: {method}")

airflow/providers/snowflake/operators/snowflake.py

@@ -20,10 +20,11 @@ from __future__ import annotations
 import time
 from collections.abc import Iterable, Mapping, Sequence
 from datetime import timedelta
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, SupportsAbs, cast
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.common.sql.operators.sql import (
     SQLCheckOperator,
     SQLExecuteQueryOperator,
@@ -34,11 +35,7 @@ from airflow.providers.snowflake.hooks.snowflake_sql_api import SnowflakeSqlApiH
 from airflow.providers.snowflake.triggers.snowflake_trigger import SnowflakeSqlApiTrigger
 
 if TYPE_CHECKING:
-    try:
-        from airflow.sdk.definitions.context import Context
-    except ImportError:
-        # TODO: Remove once provider drops support for Airflow 2
-        from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class SnowflakeCheckOperator(SQLCheckOperator):
@@ -75,8 +72,6 @@ class SnowflakeCheckOperator(SQLCheckOperator):
         Template references are recognized by str ending in '.sql'
     :param snowflake_conn_id: Reference to
         :ref:`Snowflake connection id<howto/connection:snowflake>`
-    :param autocommit: if True, each command is automatically committed.
-        (default value: True)
     :param parameters: (optional) the parameters to render the SQL query with.
     :param warehouse: name of warehouse (will overwrite any warehouse
         defined in the connection's extra JSON)
@@ -108,8 +103,6 @@ class SnowflakeCheckOperator(SQLCheckOperator):
         sql: str,
         snowflake_conn_id: str = "snowflake_default",
         parameters: Iterable | Mapping[str, Any] | None = None,
-        autocommit: bool = True,
-        do_xcom_push: bool = True,
         warehouse: str | None = None,
         database: str | None = None,
         role: str | None = None,
@@ -178,8 +171,6 @@ class SnowflakeValueCheckOperator(SQLValueCheckOperator):
         tolerance: Any = None,
         snowflake_conn_id: str = "snowflake_default",
         parameters: Iterable | Mapping[str, Any] | None = None,
-        autocommit: bool = True,
-        do_xcom_push: bool = True,
         warehouse: str | None = None,
         database: str | None = None,
         role: str | None = None,
@@ -201,7 +192,12 @@ class SnowflakeValueCheckOperator(SQLValueCheckOperator):
                 **hook_params,
             }
         super().__init__(
-            sql=sql, pass_value=pass_value, tolerance=tolerance, conn_id=snowflake_conn_id, **kwargs
+            sql=sql,
+            pass_value=pass_value,
+            tolerance=tolerance,
+            conn_id=snowflake_conn_id,
+            parameters=parameters,
+            **kwargs,
         )
         self.query_ids: list[str] = []
 
@@ -258,9 +254,6 @@ class SnowflakeIntervalCheckOperator(SQLIntervalCheckOperator):
         date_filter_column: str = "ds",
         days_back: SupportsAbs[int] = -7,
         snowflake_conn_id: str = "snowflake_default",
-        parameters: Iterable | Mapping[str, Any] | None = None,
-        autocommit: bool = True,
-        do_xcom_push: bool = True,
         warehouse: str | None = None,
         database: str | None = None,
         role: str | None = None,
@@ -354,6 +347,7 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
             When executing the statement, Snowflake replaces placeholders (? and :name) in
             the statement with these specified values.
     :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.
     """
 
     LIFETIME = timedelta(minutes=59)  # The tokens will have a 59 minutes lifetime
@@ -380,6 +374,7 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
         token_renewal_delta: timedelta = RENEWAL_DELTA,
         bindings: dict[str, Any] | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        snowflake_api_retry_args: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         self.snowflake_conn_id = snowflake_conn_id
@@ -389,7 +384,9 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
         self.token_renewal_delta = token_renewal_delta
         self.bindings = bindings
         self.execute_async = False
+        self.snowflake_api_retry_args = snowflake_api_retry_args or {}
         self.deferrable = deferrable
+        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
             kwargs["hook_params"] = {
@@ -403,6 +400,17 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
             }
         super().__init__(conn_id=snowflake_conn_id, **kwargs)  # pragma: no cover
 
+    @cached_property
+    def _hook(self):
+        return SnowflakeSqlApiHook(
+            snowflake_conn_id=self.snowflake_conn_id,
+            token_life_time=self.token_life_time,
+            token_renewal_delta=self.token_renewal_delta,
+            deferrable=self.deferrable,
+            api_retry_args=self.snowflake_api_retry_args,
+            **self.hook_params,
+        )
+
     def execute(self, context: Context) -> None:
         """
         Make a POST API request to snowflake by using SnowflakeSQL and execute the query to get the ids.
@@ -410,15 +418,8 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
         By deferring the SnowflakeSqlApiTrigger class passed along with query ids.
         """
         self.log.info("Executing: %s", self.sql)
-        self._hook = SnowflakeSqlApiHook(
-            snowflake_conn_id=self.snowflake_conn_id,
-            token_life_time=self.token_life_time,
-            token_renewal_delta=self.token_renewal_delta,
-            deferrable=self.deferrable,
-            **self.hook_params,
-        )
         self.query_ids = self._hook.execute_query(
-            self.sql,  # type: ignore[arg-type]
+            self.sql,
             statement_count=self.statement_count,
             bindings=self.bindings,
         )
@@ -504,9 +505,18 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
                 msg = f"{event['status']}: {event['message']}"
                 raise AirflowException(msg)
             if "status" in event and event["status"] == "success":
-                hook = SnowflakeSqlApiHook(snowflake_conn_id=self.snowflake_conn_id)
-                query_ids = cast("list[str]", event["statement_query_ids"])
-                hook.check_query_output(query_ids)
+                self.query_ids = cast("list[str]", event["statement_query_ids"])
+                self._hook.check_query_output(self.query_ids)
                 self.log.info("%s completed successfully.", self.task_id)
+                # Re-assign query_ids to hook after coming back from deferral to be consistent for listeners.
+                if not self._hook.query_ids:
+                    self._hook.query_ids = self.query_ids
         else:
             self.log.info("%s completed successfully.", self.task_id)
+
+    def on_kill(self) -> None:
+        """Cancel the running query."""
+        if self.query_ids:
+            self.log.info("Cancelling the query ids %s", self.query_ids)
+            self._hook.cancel_queries(self.query_ids)
+            self.log.info("Query ids %s cancelled successfully", self.query_ids)

airflow/providers/snowflake/operators/snowpark.py

@@ -17,8 +17,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Collection, Mapping, Sequence
-from typing import Any, Callable
+from collections.abc import Callable, Collection, Mapping, Sequence
+from typing import Any
 
 from airflow.providers.common.compat.standard.operators import PythonOperator, get_current_context
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook

airflow/providers/snowflake/transfers/copy_into_snowflake.py

@@ -22,11 +22,20 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import Any
 
-from airflow.models import BaseOperator
+from airflow.providers.common.compat.sdk import BaseOperator
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
 from airflow.providers.snowflake.utils.common import enclose_param
 
 
+def _validate_parameter(param_name: str, value: str | None) -> str | None:
+    """Validate that the parameter doesn't contain any invalid pattern."""
+    if value is None:
+        return None
+    if ";" in value:
+        raise ValueError(f"Invalid {param_name}: semicolons (;) not allowed.")
+    return value
+
+
 class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
     """
     Executes a COPY INTO command to load files from an external stage from clouds to Snowflake.
@@ -91,8 +100,8 @@ class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
     ):
         super().__init__(**kwargs)
         self.files = files
-        self.table = table
-        self.stage = stage
+        self.table = _validate_parameter("table", table)
+        self.stage = _validate_parameter("stage", stage)
         self.prefix = prefix
         self.file_format = file_format
         self.schema = schema
@@ -126,7 +135,7 @@ class CopyFromExternalStageToSnowflakeOperator(BaseOperator):
         if self.schema:
             into = f"{self.schema}.{self.table}"
         else:
-            into = self.table
+            into = self.table  # type: ignore[assignment]
 
         if self.columns_array:
             into = f"{into}({', '.join(self.columns_array)})"

airflow/providers/snowflake/triggers/snowflake_trigger.py

@@ -74,7 +74,6 @@ class SnowflakeSqlApiTrigger(BaseTrigger):
             self.token_renewal_delta,
         )
         try:
-            statement_query_ids: list[str] = []
             for query_id in self.query_ids:
                 while True:
                     statement_status = await self.get_query_status(query_id)
@@ -84,12 +83,10 @@ class SnowflakeSqlApiTrigger(BaseTrigger):
                 if statement_status["status"] == "error":
                     yield TriggerEvent(statement_status)
                     return
-                if statement_status["status"] == "success":
-                    statement_query_ids.extend(statement_status["statement_handles"])
             yield TriggerEvent(
                 {
                     "status": "success",
-                    "statement_query_ids": statement_query_ids,
+                    "statement_query_ids": self.query_ids,
                 }
             )
         except Exception as e: