From 8e2b757e602197d2aa5406810a8b9694008c502b Mon Sep 17 00:00:00 2001
From: Adam Ling <adam.ling@snowflake.com>
Date: Mon, 21 Oct 2024 10:00:52 -0700
Subject: [PATCH] SNOW-802436: improve error message for sql execution
 cancellation due to timeout (#2073)

---
 DESCRIPTION.md                    |  3 +++
 src/snowflake/connector/_utils.py | 11 +++++++++++
 src/snowflake/connector/cursor.py | 16 +++++++++++++---
 test/integ/test_cursor.py         |  6 +++++-
 test/unit/test_util.py            | 17 +++++++++++++++++
 5 files changed, 49 insertions(+), 4 deletions(-)
 create mode 100644 test/unit/test_util.py

diff --git a/DESCRIPTION.md b/DESCRIPTION.md
index 71d77406b..3c130ba23 100644
--- a/DESCRIPTION.md
+++ b/DESCRIPTION.md
@@ -8,6 +8,9 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 # Release Notes
 
+- v3.12.3(TBD)
+  - Improved error message for SQL execution cancellations caused by timeout.
+
 - v3.12.2(September 11,2024)
   - Improved error handling for asynchronous queries, providing more detailed and informative error messages when an async query fails.
   - Improved inference of top-level domains for accounts specifying a region in China, now defaulting to snowflakecomputing.cn.
diff --git a/src/snowflake/connector/_utils.py b/src/snowflake/connector/_utils.py
index de9ea78ff..85ea83073 100644
--- a/src/snowflake/connector/_utils.py
+++ b/src/snowflake/connector/_utils.py
@@ -7,6 +7,7 @@
 import string
 from enum import Enum
 from random import choice
+from threading import Timer
 
 
 class TempObjectType(Enum):
@@ -43,3 +44,13 @@ def random_name_for_temp_object(object_type: TempObjectType) -> str:
 
 def get_temp_type_for_object(use_scoped_temp_objects: bool) -> str:
     return SCOPED_TEMPORARY_STRING if use_scoped_temp_objects else TEMPORARY_STRING
+
+
+class _TrackedQueryCancellationTimer(Timer):
+    def __init__(self, interval, function, args=None, kwargs=None):
+        super().__init__(interval, function, args, kwargs)
+        self.executed = False
+
+    def run(self):
+        super().run()
+        self.executed = True
diff --git a/src/snowflake/connector/cursor.py b/src/snowflake/connector/cursor.py
index 97bbc672e..8b9d400e0 100644
--- a/src/snowflake/connector/cursor.py
+++ b/src/snowflake/connector/cursor.py
@@ -16,7 +16,7 @@
 import warnings
 from enum import Enum
 from logging import getLogger
-from threading import Lock, Timer
+from threading import Lock
 from types import TracebackType
 from typing import (
     IO,
@@ -39,6 +39,7 @@
 
 from . import compat
 from ._sql_util import get_file_transfer_type
+from ._utils import _TrackedQueryCancellationTimer
 from .bind_upload_agent import BindUploadAgent, BindUploadError
 from .constants import (
     FIELD_NAME_TO_ID,
@@ -392,7 +393,9 @@ def __init__(
         self.messages: list[
             tuple[type[Error] | type[Exception], dict[str, str | bool]]
         ] = []
-        self._timebomb: Timer | None = None  # must be here for abort_exit method
+        self._timebomb: _TrackedQueryCancellationTimer | None = (
+            None  # must be here for abort_exit method
+        )
         self._description: list[ResultMetadataV2] | None = None
         self._sfqid: str | None = None
         self._sqlstate = None
@@ -654,7 +657,9 @@ def _execute_helper(
         )
 
         if real_timeout is not None:
-            self._timebomb = Timer(real_timeout, self.__cancel_query, [query])
+            self._timebomb = _TrackedQueryCancellationTimer(
+                real_timeout, self.__cancel_query, [query]
+            )
             self._timebomb.start()
             logger.debug("started timebomb in %ss", real_timeout)
         else:
@@ -1071,6 +1076,11 @@ def execute(
             logger.debug(ret)
             err = ret["message"]
             code = ret.get("code", -1)
+            if self._timebomb and self._timebomb.executed:
+                err = (
+                    f"SQL execution was cancelled by the client due to a timeout. "
+                    f"Error message received from the server: {err}"
+                )
             if "data" in ret:
                 err += ret["data"].get("errorMessage", "")
             errvalue = {
diff --git a/test/integ/test_cursor.py b/test/integ/test_cursor.py
index 7d6f82570..384e5e95a 100644
--- a/test/integ/test_cursor.py
+++ b/test/integ/test_cursor.py
@@ -767,7 +767,11 @@ def test_timeout_query(conn_cnx):
                     "select seq8() as c1 from table(generator(timeLimit => 60))",
                     timeout=5,
                 )
-            assert err.value.errno == 604, "Invalid error code"
+            assert err.value.errno == 604, (
+                "Invalid error code"
+                and "SQL execution was cancelled by the client due to a timeout"
+                in err.value.msg
+            )
 
 
 def test_executemany(conn, db_parameters):
diff --git a/test/unit/test_util.py b/test/unit/test_util.py
new file mode 100644
index 000000000..9fa445f5d
--- /dev/null
+++ b/test/unit/test_util.py
@@ -0,0 +1,17 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from snowflake.connector._utils import _TrackedQueryCancellationTimer
+
+
+def test_timer():
+    timer = _TrackedQueryCancellationTimer(1, lambda: None)
+    timer.start()
+    timer.join()
+    assert timer.executed
+
+    timer = _TrackedQueryCancellationTimer(1, lambda: None)
+    timer.start()
+    timer.cancel()
+    assert not timer.executed