misc fixes for psycopg instrumentation

Author: sohankshirsagar

drift/instrumentation/psycopg/instrumentation.py

@@ -483,9 +483,11 @@ def _traced_executemany(
             # RECORD mode: Execute real query and record span
             time.time()
             error = None
+            # Convert to list BEFORE executing to avoid iterator exhaustion
+            params_list = list(params_seq)
 
             try:
-                result = original_executemany(query, params_seq, **kwargs)
+                result = original_executemany(query, params_list, **kwargs)
                 return result
             except Exception as e:
                 error = e
@@ -494,7 +496,6 @@ def _traced_executemany(
                 # Always create span in RECORD mode (including pre-app-start queries)
                 # Pre-app-start queries are marked with is_pre_app_start=true flag
                 if sdk.mode == TuskDriftMode.RECORD:
-                    params_list = list(params_seq)
                     self._finalize_query_span(
                         span,
                         cursor,

drift/instrumentation/psycopg2/instrumentation.py

@@ -211,13 +211,17 @@ def __init__(self, connection: Any, instrumentation: Psycopg2Instrumentation, sd
         object.__setattr__(self, "_connection", connection)
         object.__setattr__(self, "_instrumentation", instrumentation)
         object.__setattr__(self, "_sdk", sdk)
+        # Preserve the connection's default cursor_factory (set at connect() time)
+        object.__setattr__(self, "_default_cursor_factory", getattr(connection, "cursor_factory", None))
 
     def cursor(self, name: str | None = None, cursor_factory: Any = None, *args: Any, **kwargs: Any) -> Any:
         """Intercept cursor creation to wrap user-provided cursor_factory."""
-        # Create instrumented cursor factory (wrapping user's factory if provided)
+        # Use cursor_factory from cursor() call, or fall back to connection's default
+        base_factory = cursor_factory if cursor_factory is not None else self._default_cursor_factory
+        # Create instrumented cursor factory (wrapping the base factory)
         wrapped_factory = self._instrumentation._create_cursor_factory(
             self._sdk,
-            cursor_factory,  # This becomes the base class (None uses default)
+            base_factory,
         )
         return self._connection.cursor(*args, name=name, cursor_factory=wrapped_factory, **kwargs)
 
@@ -606,11 +610,12 @@ def _traced_executemany(
 
                 # For all other queries (pre-app-start OR within a request trace), get mock
                 # Wrap in {"_batch": ...} to match the recording format
+                # Normalize to list to match RECORD mode and avoid iterator/serialization issues
                 is_pre_app_start = not sdk.app_ready
                 mock_result = self._try_get_mock(
                     sdk,
                     query,
-                    {"_batch": params_list},
+                    {"_batch": list(params_list)},
                     trace_id,
                     span_id,
                     parent_span_id,
@@ -642,9 +647,11 @@ def _traced_executemany(
 
             # RECORD mode: Execute real query and record span
             error = None
+            # Convert to list BEFORE executing to avoid iterator exhaustion
+            params_as_list = list(params_list)
 
             try:
-                result = original_executemany(query, params_list)
+                result = original_executemany(query, params_as_list)
                 return result
             except Exception as e:
                 error = e
@@ -656,7 +663,7 @@ def _traced_executemany(
                         span,
                         cursor,
                         query,
-                        {"_batch": list(params_list)},
+                        {"_batch": params_as_list},
                         error,
                     )
         finally:

drift/instrumentation/utils/psycopg_utils.py

@@ -13,20 +13,26 @@ def deserialize_db_value(val: Any) -> Any:
     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).
 
+    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.
+
     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.
     """
     if isinstance(val, str):
-        # Try to parse as ISO datetime
-        try:
-            # Handle Z suffix for UTC
-            parsed = dt.datetime.fromisoformat(val.replace("Z", "+00:00"))
-            return parsed
-        except ValueError:
-            pass
+        # 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
+        if ("T" in val or (" " in val and ":" in val)) and "-" in val:
+            try:
+                # Handle Z suffix for UTC
+                parsed = dt.datetime.fromisoformat(val.replace("Z", "+00:00"))
+                return parsed
+            except ValueError:
+                pass
     elif isinstance(val, list):
         return [deserialize_db_value(v) for v in val]
     elif isinstance(val, dict):