handle replaying uuid properly

Author: sohankshirsagar

drift/instrumentation/utils/psycopg_utils.py

@@ -4,6 +4,7 @@
 
 import base64
 import datetime as dt
+import uuid
 from typing import Any
 
 
@@ -13,6 +14,7 @@ def deserialize_db_value(val: Any) -> Any:
     During recording, database values are serialized for JSON storage:
     - datetime objects -> ISO format strings
     - bytes/memoryview -> {"__bytes__": "<base64_encoded_data>"}
+    - uuid.UUID -> {"__uuid__": "<uuid_string>"}
 
     During replay, we need to convert them back to their original types so that
     application code (Flask/Django) handles them the same way.
@@ -28,6 +30,9 @@ def deserialize_db_value(val: Any) -> Any:
         if "__bytes__" in val and len(val) == 1:
             # Decode base64 back to bytes
             return base64.b64decode(val["__bytes__"])
+        # Check for UUID tagged structure
+        if "__uuid__" in val and len(val) == 1:
+            return uuid.UUID(val["__uuid__"])
         # Recursively deserialize dict values
         return {k: deserialize_db_value(v) for k, v in val.items()}
     elif isinstance(val, str):

drift/instrumentation/utils/serialization.py

@@ -43,7 +43,7 @@ def serialize_value(val: Any) -> Any:
     if isinstance(val, (datetime.datetime, datetime.date, datetime.time)):
         return val.isoformat()
     elif isinstance(val, uuid.UUID):
-        return str(val)
+        return {"__uuid__": str(val)}
     elif isinstance(val, memoryview):
         # Convert memoryview to bytes first, then serialize
         return _serialize_bytes(bytes(val))