allow for UUID encoding

Author: sfc-gh-lspiegelberg

src/snowflake/connector/network.py

@@ -352,6 +352,13 @@ def close(self) -> None:
         self._active_sessions.clear()
         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 +506,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 +568,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 +666,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,6 +1,10 @@
 #!/usr/bin/env python
 import io
+import json
 import unittest.mock
+import uuid
+
+from src.snowflake.connector.network import SnowflakeRestfulJsonEncoder
 from test.unit.mock_utils import mock_connection
 
 import pytest
@@ -63,3 +67,14 @@ 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