SNOW-2011670 Allow url parameter requestId to be set with statement params (#2240)

As part of the Snowpark IR project, we will send AST as part of the request. To restore them faithfully on the server-side and avoid introducing yet another layer of indirection, we would like to use the HTTP Request Id. Currently, this Request Id is generated on the client as part of issuing the query. With this PR we introduce a new (client-side) statement parameter requestId that allows to overwrite the automatically generated request Id. requestId must be a valid UUID4 or UUID4 string.

(cherry picked from commit 6e214d0)

Author: sfc-gh-lspiegelberg

DESCRIPTION.md

@@ -19,6 +19,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Fixed a bug that caused driver to fail silently on `TO_DATE` arrow to python conversion when invalid date was followed by the correct one.
   - Added `check_arrow_conversion_error_on_every_column` connection property that can be set to `False` to restore previous behaviour in which driver will ignore errors until it occurs in the last column. This flag's purpose is to unblock workflows that may be impacted by the bugfix and will be removed in later releases.
   - Lower log levels from info to debug for some of the messages to make the output easier to follow.
+  - Allow the connector to inherit a UUID4 generated upstream, provided in statement parameters (field: `requestId`), rather than automatically generate a UUID4 to use for the HTTP Request ID.
 
 - v3.14.0(March 03, 2025)
   - Bumped pyOpenSSL dependency upper boundary from <25.0.0 to <26.0.0.

src/snowflake/connector/_utils.py

@@ -4,6 +4,7 @@
 from enum import Enum
 from random import choice
 from threading import Timer
+from uuid import UUID
 
 
 class TempObjectType(Enum):
@@ -29,6 +30,8 @@ class TempObjectType(Enum):
     "PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS"
 )
 
+REQUEST_ID_STATEMENT_PARAM_NAME = "requestId"
+
 
 def generate_random_alphanumeric(length: int = 10) -> str:
     return "".join(choice(ALPHANUMERIC) for _ in range(length))
@@ -42,6 +45,21 @@ def get_temp_type_for_object(use_scoped_temp_objects: bool) -> str:
     return SCOPED_TEMPORARY_STRING if use_scoped_temp_objects else TEMPORARY_STRING
 
 
+def is_uuid4(str_or_uuid: str | UUID) -> bool:
+    """Check whether provided string str is a valid UUID version4."""
+    if isinstance(str_or_uuid, UUID):
+        return str_or_uuid.version == 4
+
+    if not isinstance(str_or_uuid, str):
+        return False
+
+    try:
+        uuid_str = str(UUID(str_or_uuid, version=4))
+    except ValueError:
+        return False
+    return uuid_str == str_or_uuid
+
+
 class _TrackedQueryCancellationTimer(Timer):
     def __init__(self, interval, function, args=None, kwargs=None):
         super().__init__(interval, function, args, kwargs)

src/snowflake/connector/cursor.py

@@ -35,7 +35,11 @@
 
 from . import compat
 from ._sql_util import get_file_transfer_type
-from ._utils import _TrackedQueryCancellationTimer
+from ._utils import (
+    REQUEST_ID_STATEMENT_PARAM_NAME,
+    _TrackedQueryCancellationTimer,
+    is_uuid4,
+)
 from .bind_upload_agent import BindUploadAgent, BindUploadError
 from .constants import (
     CMD_TYPE_DOWNLOAD,
@@ -637,7 +641,27 @@ def _execute_helper(
                     )
 
         self._sequence_counter = 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/test_cursor.py

@@ -7,6 +7,7 @@
 import os
 import pickle
 import time
+import uuid
 from datetime import date, datetime, timezone
 from typing import TYPE_CHECKING, NamedTuple
 from unittest import mock
@@ -1964,3 +1965,37 @@ def test_nanoarrow_usage_deprecation():
             and "snowflake.connector.cursor.NanoarrowUsage has been deprecated"
             in str(record[2].message)
         )
+
+
+@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"),
+    ],
+)
+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"):
+        with conn_cnx() as con:
+            with con.cursor() as cur:
+                cur.execute(
+                    "select seq4() as foo from table(generator(rowcount=>5))",
+                    _statement_params={"requestId": request_id},
+                )
+
+
+def test_custom_request_id(conn_cnx):
+    request_id = uuid.uuid4()
+
+    with conn_cnx() as con:
+        with con.cursor() as cur:
+            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."