SNOW-2027116 Allow for UUID encoding in SnowflakeRestful interface (#2254)

Author: sfc-gh-lspiegelberg

src/snowflake/connector/network.py

@@ -353,6 +353,15 @@ def close(self) -> None:
         self._idle_sessions.clear()
 
 
+# Customizable JSONEncoder to support additional types.
+class SnowflakeRestfulJsonEncoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, uuid.UUID):
+            return str(o)
+
+        return super().default(o)
+
+
 class SnowflakeRestful:
     """Snowflake Restful class."""
 
@@ -499,7 +508,7 @@ def request(
             return self._post_request(
                 url,
                 headers,
-                json.dumps(body),
+                json.dumps(body, cls=SnowflakeRestfulJsonEncoder),
                 token=self.token,
                 _no_results=_no_results,
                 timeout=timeout,
@@ -561,7 +570,7 @@ def _token_request(self, request_type):
         ret = self._post_request(
             url,
             headers,
-            json.dumps(body),
+            json.dumps(body, cls=SnowflakeRestfulJsonEncoder),
             token=header_token,
         )
         if ret.get("success") and ret.get("data", {}).get("sessionToken"):
@@ -659,7 +668,7 @@ def delete_session(self, retry: bool = False) -> None:
                 ret = self._post_request(
                     url,
                     headers,
-                    json.dumps(body),
+                    json.dumps(body, cls=SnowflakeRestfulJsonEncoder),
                     token=self.token,
                     timeout=5,
                     no_retry=True,

test/unit/test_network.py

@@ -1,10 +1,14 @@
 #!/usr/bin/env python
 import io
+import json
 import unittest.mock
+import uuid
 from test.unit.mock_utils import mock_connection
 
 import pytest
 
+from src.snowflake.connector.network import SnowflakeRestfulJsonEncoder
+
 try:
     from snowflake.connector import Error, InterfaceError
     from snowflake.connector.network import SnowflakeRestful
@@ -63,3 +67,20 @@ def test_fetch():
         # if no retry is set to False, the function raises an InterfaceError
         with pytest.raises(InterfaceError) as exc:
             assert rest.fetch(**default_parameters, no_retry=False)
+
+
+@pytest.mark.parametrize(
+    "u",
+    [
+        uuid.uuid1(),
+        uuid.uuid3(uuid.NAMESPACE_URL, "www.snowflake.com"),
+        uuid.uuid4(),
+        uuid.uuid5(uuid.NAMESPACE_URL, "www.snowflake.com"),
+    ],
+)
+def test_json_serialize_uuid(u):
+    expected = f'{{"u": "{u}", "a": 42}}'
+
+    assert (json.dumps(u, cls=SnowflakeRestfulJsonEncoder)) == f'"{u}"'
+
+    assert json.dumps({"u": u, "a": 42}, cls=SnowflakeRestfulJsonEncoder) == expected