[ASYNC] Apply #2236 to async code

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_cursor.py

@@ -685,7 +685,15 @@ async def execute(
             logger.debug(ret)
             err = ret["message"]
             code = ret.get("code", -1)
-            if self._timebomb and self._timebomb.result():
+            if (
+                self._timebomb
+                and self._timebomb.result()
+                and "SQL execution canceled" in err
+            ):
+                # Modify the error message only if the server error response indicates the query was canceled.
+                # If the error occurs before the cancellation request reaches the backend
+                # (e.g., due to a very short timeout), we retain the original error message
+                # as the query might have encountered an issue prior to cancellation.
                 err = (
                     f"SQL execution was cancelled by the client due to a timeout. "
                     f"Error message received from the server: {err}"

test/integ/aio/test_cursor_async.py

@@ -15,6 +15,7 @@
 from datetime import date, datetime, timezone
 from typing import NamedTuple
 from unittest import mock
+from unittest.mock import MagicMock
 
 import pytest
 import pytz
@@ -792,6 +793,7 @@ async def test_invalid_bind_data_type(conn_cnx):
             await cnx.cursor().execute("select 1 from dual where 1=%s", ([1, 2, 3],))
 
 
+@pytest.mark.skipolddriver
 async def test_timeout_query(conn_cnx):
     async with conn_cnx() as cnx:
         async with cnx.cursor() as c:
@@ -802,10 +804,32 @@ async def test_timeout_query(conn_cnx):
                 )
             assert err.value.errno == 604, (
                 "Invalid error code"
-                and "SQL execution was cancelled by the client due to a timeout"
+                and "SQL execution was cancelled by the client due to a timeout. Error message received from the server: SQL execution canceled"
                 in err.value.msg
             )
 
+            with pytest.raises(errors.ProgrammingError) as err:
+                # we can not precisely control the timing to send cancel query request right after server
+                # executes the query but before returning the results back to client
+                # it depends on python scheduling and server processing speed, so we mock here
+                mock_timebomb = MagicMock()
+                mock_timebomb.result.return_value = True
+
+                with mock.patch.object(c, "_timebomb", mock_timebomb):
+                    await c.execute(
+                        "select 123'",
+                        timeout=0.1,
+                    )
+            assert (
+                mock_timebomb.result.return_value is True and err.value.errno == 1003
+            ), (
+                "Invalid error code"
+                and "SQL compilation error:\nsyntax error line 1 at position 10 unexpected '''."
+                in err.value.msg
+                and "SQL execution was cancelled by the client due to a timeout"
+                not in err.value.msg
+            )
+
 
 async def test_executemany(conn, db_parameters):
     """Executes many statements. Client binding is supported by either dict, or list data types.