[Async] Apply #2240 to async code

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/_cursor.py

@@ -56,6 +56,8 @@
 from snowflake.connector.telemetry import TelemetryData, TelemetryField
 from snowflake.connector.time_util import get_time_millis
 
+from .._utils import REQUEST_ID_STATEMENT_PARAM_NAME, is_uuid4
+
 if TYPE_CHECKING:
     from pandas import DataFrame
     from pyarrow import Table
@@ -202,7 +204,27 @@ async def _execute_helper(
                     )
 
         self._sequence_counter = await self._connection._next_sequence_counter()
-        self._request_id = uuid.uuid4()
+
+        # If requestId is contained in statement parameters, use it to set request id. Verify here it is a valid uuid4
+        # identifier.
+        if (
+            statement_params is not None
+            and REQUEST_ID_STATEMENT_PARAM_NAME in statement_params
+        ):
+            request_id = statement_params[REQUEST_ID_STATEMENT_PARAM_NAME]
+
+            if not is_uuid4(request_id):
+                # uuid.UUID will throw an error if invalid, but we explicitly check and throw here.
+                raise ValueError(f"requestId {request_id} is not a valid UUID4.")
+            self._request_id = uuid.UUID(str(request_id), version=4)
+
+            # Create a (deep copy) and remove the statement param, there is no need to encode it as extra parameter
+            # one more time.
+            statement_params = statement_params.copy()
+            statement_params.pop(REQUEST_ID_STATEMENT_PARAM_NAME)
+        else:
+            # Generate UUID for query.
+            self._request_id = uuid.uuid4()
 
         logger.debug(f"Request id: {self._request_id}")
 

test/integ/aio/test_cursor_async.py

@@ -12,6 +12,7 @@
 import os
 import pickle
 import time
+import uuid
 from datetime import date, datetime, timezone
 from typing import NamedTuple
 from unittest import mock
@@ -1868,3 +1869,37 @@ async def test_fetch_download_timeout_setting(conn_cnx):
         sql = "SELECT seq4(), uniform(1, 10, RANDOM(12)) FROM TABLE(GENERATOR(ROWCOUNT => 100000)) v"
         async with conn_cnx() as con, con.cursor() as cur:
             assert len(await (await cur.execute(sql)).fetchall()) == 100000
+
+
+@pytest.mark.parametrize(
+    "request_id",
+    [
+        "THIS IS NOT VALID",
+        uuid.uuid1(),
+        uuid.uuid3(uuid.NAMESPACE_URL, "www.snowflake.com"),
+        uuid.uuid5(uuid.NAMESPACE_URL, "www.snowflake.com"),
+    ],
+)
+async def test_custom_request_id_negative(request_id, conn_cnx):
+
+    # Ensure that invalid request_ids (non uuid4) do not compromise interface.
+    with pytest.raises(ValueError, match="requestId"):
+        async with conn_cnx() as con:
+            async with con.cursor() as cur:
+                await cur.execute(
+                    "select seq4() as foo from table(generator(rowcount=>5))",
+                    _statement_params={"requestId": request_id},
+                )
+
+
+async def test_custom_request_id(conn_cnx):
+    request_id = uuid.uuid4()
+
+    async with conn_cnx() as con:
+        async with con.cursor() as cur:
+            await cur.execute(
+                "select seq4() as foo from table(generator(rowcount=>5))",
+                _statement_params={"requestId": request_id},
+            )
+
+            assert cur._sfqid is not None, "Query must execute successfully."