From 1efe40b2881177c8975fca16b0791bd9ccf9fbd7 Mon Sep 17 00:00:00 2001
From: Ryan Tan <65996005+ryn5@users.noreply.github.com>
Date: Thu, 29 Feb 2024 08:56:10 -0800
Subject: [PATCH] TINKERPOP-3054 Fix requestId Deserialization in
 `gremlin-python` (#2494)

https://issues.apache.org/jira/browse/TINKERPOP-3054
Added str() guards to request_id before validating to avoid passing UUID into UUID constructor, which is intended to receive a string.
---
 CHANGELOG.asciidoc                                         | 5 +++--
 .../src/main/python/gremlin_python/driver/connection.py    | 7 ++++---
 .../src/main/python/gremlin_python/driver/serializer.py    | 2 +-
 gremlin-python/src/main/python/tests/driver/test_client.py | 6 +++++-
 4 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.asciidoc b/CHANGELOG.asciidoc
index b70b64ccd7d..6f2e3d63c31 100644
--- a/CHANGELOG.asciidoc
+++ b/CHANGELOG.asciidoc
@@ -24,14 +24,15 @@ image::https://raw.githubusercontent.com/apache/tinkerpop/master/docs/static/ima
 === TinkerPop 3.6.7 (NOT OFFICIALLY RELEASED YET)
 
 * Fixed a bug in Gremlin.Net for .NET 8 that led to exceptions: `InvalidOperationException: Enumeration has not started. Call MoveNext.`
+* Fixed message requestId serialization in `gremlin-python`.
 * Fixed bug in bytecode translation of `g.tx().commit()` and `g.tx().rollback()` in all languages.
 * Improved error message from `JavaTranslator` by including exception source.
-* Added missing `short` serialization (`gx:Int16`) to GraphSONV2 and GraphSONV3 in `gremlin-python`
+* Added missing `short` serialization (`gx:Int16`) to GraphSONV2 and GraphSONV3 in `gremlin-python`.
 * Added tests for error handling for GLV's if tx.commit() is called remotely for graphs without transactions support.
 * Introduced multi-architecture AMD64/ARM64 docker images for gremlin-console.
 * Fixed bug in `JavaTranslator` where `has(String, null)` could call `has(String, Traversal)` to generate an error.
 * Fixed issue where server errors weren't being properly parsed when sending bytecode over HTTP.
-* Improved bulkset contains check for elements if all elements in bulkset are of the same type
+* Improved bulkset contains check for elements if all elements in bulkset are of the same type.
 * Fixed bug in `EarlyLimitStrategy` which was too aggressive when promoting `limit()` before `map()`.
 * Fixed bug in mid-traversal `mergeE()` where mutations in `sideEffect()` were being applied to the current traverser rather than a `onMatch` edge.
 
diff --git a/gremlin-python/src/main/python/gremlin_python/driver/connection.py b/gremlin-python/src/main/python/gremlin_python/driver/connection.py
index 9641e4e203a..789de181d02 100644
--- a/gremlin-python/src/main/python/gremlin_python/driver/connection.py
+++ b/gremlin-python/src/main/python/gremlin_python/driver/connection.py
@@ -58,10 +58,11 @@ def close(self):
     def write(self, request_message):
         if not self._inited:
             self.connect()
-        request_id = str(uuid.uuid4())
         if request_message.args.get("requestId"):
-            request_id = request_message.args.get("requestId")
-            uuid.UUID(request_id) # Checks for proper UUID or else server will return an error.
+            request_id = str(request_message.args.get("requestId"))
+            uuid.UUID(request_id)  # Checks for proper UUID or else server will return an error.
+        else:
+            request_id = str(uuid.uuid4())
         result_set = resultset.ResultSet(queue.Queue(), request_id)
         self._results[request_id] = result_set
         # Create write task
diff --git a/gremlin-python/src/main/python/gremlin_python/driver/serializer.py b/gremlin-python/src/main/python/gremlin_python/driver/serializer.py
index 74ed99523c2..6b522e50b39 100644
--- a/gremlin-python/src/main/python/gremlin_python/driver/serializer.py
+++ b/gremlin-python/src/main/python/gremlin_python/driver/serializer.py
@@ -240,7 +240,7 @@ def build_message(self, request_id, processor, op, args):
     def finalize_message(self, message, mime_len, mime_type):
         ba = bytearray()
 
-        request_id = uuid.UUID(message['requestId'])
+        request_id = uuid.UUID(str(message['requestId']))
         ba.extend(self.header_pack(mime_len, mime_type, 0x81,
                                    (request_id.int >> 64) & self.max_int64, request_id.int & self.max_int64))
 
diff --git a/gremlin-python/src/main/python/tests/driver/test_client.py b/gremlin-python/src/main/python/tests/driver/test_client.py
index 221b825ea6d..9bff038eb18 100644
--- a/gremlin-python/src/main/python/tests/driver/test_client.py
+++ b/gremlin-python/src/main/python/tests/driver/test_client.py
@@ -453,7 +453,11 @@ def test_client_custom_invalid_request_id_graphbinary_script(client):
         assert "badly formed hexadecimal UUID string" in str(ex)
 
 
-def test_client_custom_valid_request_id_script(client):
+def test_client_custom_valid_request_id_script_uuid(client):
+    assert len(client.submit('g.V()', request_options={"requestId":uuid.uuid4()}).all().result()) == 6
+
+
+def test_client_custom_valid_request_id_script_string(client):
     assert len(client.submit('g.V()', request_options={"requestId":str(uuid.uuid4())}).all().result()) == 6