fix bytea serialization to preserve binary data during record/replay

Author: sohankshirsagar

drift/instrumentation/psycopg/e2e-tests/src/app.py

@@ -559,19 +559,11 @@ def test_sql_composed():
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-
-
-
-# ===== BUG HUNTING TEST ENDPOINTS =====
-# These endpoints expose confirmed bugs in the psycopg instrumentation
-# Endpoints that passed tests have been removed
-
 @app.route("/test/binary-uuid")
 def test_binary_uuid():
     """Test binary UUID data type.
 
-    BUG HYPOTHESIS: UUID types may not serialize/deserialize correctly
-    during RECORD/REPLAY because they are binary.
+    Tests whether the instrumentation correctly handles binary UUID data types.
     """
     try:
         import uuid
@@ -602,13 +594,11 @@ def test_binary_uuid():
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-
 @app.route("/test/binary-bytea")
 def test_binary_bytea():
     """Test binary bytea data type.
 
-    BUG HYPOTHESIS: Binary data (bytea) may not serialize/deserialize
-    correctly during RECORD/REPLAY.
+    Tests whether the instrumentation correctly handles binary bytea data types.
     """
     try:
         with psycopg.connect(get_conn_string()) as conn, conn.cursor() as cur:

drift/instrumentation/psycopg/e2e-tests/src/test_requests.py

@@ -96,14 +96,8 @@ def make_request(method, endpoint, **kwargs):
 
     make_request("GET", "/test/sql-composed")
 
-    # ===== BUG HUNTING TEST ENDPOINTS =====
-    # These tests expose confirmed bugs in the psycopg instrumentation
-    # See BUG_TRACKING.md for detailed information about each bug
-    print("\n--- Bug Hunting Tests (REPLAY mode bugs - pass RECORD but fail REPLAY) ---\n")
-
-    # Bug 8: UUID parameter serialization issue during REPLAY
     make_request("GET", "/test/binary-uuid")
-    # Bug 9: bytea data deserialization returns string instead of bytes
+
     make_request("GET", "/test/binary-bytea")
 
     print("\nAll requests completed successfully")

drift/instrumentation/utils/psycopg_utils.py

@@ -2,27 +2,35 @@
 
 from __future__ import annotations
 
+import base64
 import datetime as dt
 from typing import Any
 
 
 def deserialize_db_value(val: Any) -> Any:
-    """Convert ISO datetime strings back to datetime objects for consistent serialization.
+    """Convert serialized values back to their original Python types.
 
-    During recording, datetime objects from the database are serialized to ISO format strings.
-    During replay, we need to convert them back to datetime objects so that Flask/Django
-    serializes them the same way (e.g., RFC 2822 vs ISO 8601 format).
+    During recording, database values are serialized for JSON storage:
+    - datetime objects -> ISO format strings
+    - bytes/memoryview -> {"__bytes__": "<base64_encoded_data>"}
 
-    Only parses strings that contain a time component (T or space separator with :) to avoid
-    incorrectly converting date-only strings or text that happens to look like dates.
+    During replay, we need to convert them back to their original types so that
+    application code (Flask/Django) handles them the same way.
 
     Args:
         val: A value from the mocked database rows. Can be a string, list, dict, or any other type.
 
     Returns:
-        The value with ISO datetime strings converted back to datetime objects.
+        The value with serialized types converted back to their original Python types.
     """
-    if isinstance(val, str):
+    if isinstance(val, dict):
+        # Check for bytes tagged structure
+        if "__bytes__" in val and len(val) == 1:
+            # Decode base64 back to bytes
+            return base64.b64decode(val["__bytes__"])
+        # Recursively deserialize dict values
+        return {k: deserialize_db_value(v) for k, v in val.items()}
+    elif isinstance(val, str):
         # Only parse strings that look like full datetime (must have time component)
         # This avoids converting date-only strings like "2024-01-15" or text columns
         # that happen to match date patterns
@@ -35,6 +43,4 @@ def deserialize_db_value(val: Any) -> Any:
                 pass
     elif isinstance(val, list):
         return [deserialize_db_value(v) for v in val]
-    elif isinstance(val, dict):
-        return {k: deserialize_db_value(v) for k, v in val.items()}
     return val

drift/instrumentation/utils/serialization.py

@@ -2,11 +2,33 @@
 
 from __future__ import annotations
 
+import base64
 import datetime
 import uuid
 from typing import Any
 
 
+def _serialize_bytes(val: bytes) -> Any:
+    """Serialize bytes to a JSON-compatible format.
+
+    Attempts UTF-8 decode first for text data (like COPY output).
+    Falls back to base64 encoding with tagged structure for binary data
+    that contains invalid UTF-8 sequences (like bytea columns).
+
+    Args:
+        val: The bytes value to serialize.
+
+    Returns:
+        Either a string (if valid UTF-8) or a dict {"__bytes__": "base64_data"}.
+    """
+    try:
+        # Try UTF-8 decode first - works for text data like COPY output
+        return val.decode("utf-8")
+    except UnicodeDecodeError:
+        # Fall back to base64 for binary data with invalid UTF-8 sequences
+        return {"__bytes__": base64.b64encode(val).decode("ascii")}
+
+
 def serialize_value(val: Any) -> Any:
     """Convert non-JSON-serializable values to JSON-compatible types.
 
@@ -22,10 +44,11 @@ def serialize_value(val: Any) -> Any:
         return val.isoformat()
     elif isinstance(val, uuid.UUID):
         return str(val)
-    elif isinstance(val, bytes):
-        return val.decode("utf-8", errors="replace")
     elif isinstance(val, memoryview):
-        return bytes(val).decode("utf-8", errors="replace")
+        # Convert memoryview to bytes first, then serialize
+        return _serialize_bytes(bytes(val))
+    elif isinstance(val, bytes):
+        return _serialize_bytes(val)
     elif isinstance(val, (list, tuple)):
         return [serialize_value(v) for v in val]
     elif isinstance(val, dict):