fix(FIR-43199): add query cancel setting

src/firebolt/async_db/cursor.py

@@ -263,6 +263,9 @@ async def _handle_query_execution(
         query_params: Dict[str, Any] = {"output_format": JSON_OUTPUT_FORMAT}
         if async_execution:
             query_params["async"] = True
+        # timeout cancels the query by aborting the connection
+        if timeout_controller.timeout is not None:
+            query_params["cancel_query_on_connection_drop"] = "all"
         resp = await self._api_request(
             query,
             query_params,

src/firebolt/db/cursor.py

@@ -257,6 +257,9 @@ def _handle_query_execution(
         query_params: Dict[str, Any] = {"output_format": JSON_OUTPUT_FORMAT}
         if async_execution:
             query_params["async"] = True
+        # timeout cancels the query by aborting the connection
+        if timeout_controller.timeout is not None:
+            query_params["cancel_query_on_connection_drop"] = "all"
         resp = self._api_request(
             query,
             query_params,

tests/integration/dbapi/async/V2/test_timeout.py

@@ -1,3 +1,5 @@
+import time
+
 from pytest import raises
 
 from firebolt.async_db import Connection
@@ -8,6 +10,17 @@
 
 
 async def test_query_timeout(connection: Connection):
+    timestamp = int(time.time())
+    label = f"test_query_async_timeout_cancel:{timestamp}"
     with connection.cursor() as cursor:
         with raises(QueryTimeoutError):
+            await cursor.execute(f"SET query_label='{label}'")
             await cursor.execute(LONG_SELECT, timeout_seconds=1)
+        time.sleep(10)  # it takes some time for query history to update
+        await cursor.execute("SET query_label=''")
+        await cursor.execute(
+            "SELECT status FROM information_schema.engine_query_history WHERE query_label = ? ORDER BY end_time DESC",
+            [label],
+        )
+        status = await cursor.fetchone()
+        assert status[0] == "CANCELED_EXECUTION"

tests/integration/dbapi/sync/V2/test_timeout.py

@@ -1,3 +1,5 @@
+import time
+
 from pytest import raises
 
 from firebolt.db import Connection
@@ -7,6 +9,17 @@
 
 
 def test_query_timeout(connection: Connection):
+    timestamp = int(time.time())
+    label = f"test_query_async_timeout_cancel:{timestamp}"
     with connection.cursor() as cursor:
         with raises(QueryTimeoutError):
+            cursor.execute(f"SET query_label='{label}'")
             cursor.execute(LONG_SELECT, timeout_seconds=1)
+        time.sleep(10)  # it takes some time for query history to update
+        cursor.execute("SET query_label=''")
+        cursor.execute(
+            "SELECT status FROM information_schema.engine_query_history WHERE query_label = ? ORDER BY end_time DESC",
+            [label],
+        )
+        status = cursor.fetchone()
+        assert status[0] == "CANCELED_EXECUTION"