refactored psycopg and psycopg2

Author: sohankshirsagar

drift/instrumentation/psycopg/instrumentation.py

@@ -6,16 +6,17 @@
 from types import ModuleType
 from typing import Any
 
-from opentelemetry import context as otel_context
 from opentelemetry import trace
 from opentelemetry.trace import SpanKind as OTelSpanKind
-from opentelemetry.trace import Status, set_span_in_context
+from opentelemetry.trace import Status
 from opentelemetry.trace import StatusCode as OTelStatusCode
 
 from ...core.communication.types import MockRequestInput
 from ...core.drift_sdk import TuskDrift
 from ...core.json_schema_helper import JsonSchemaHelper
+from ...core.mode_utils import handle_record_mode, handle_replay_mode
 from ...core.tracing import TdSpanAttributes
+from ...core.tracing.span_utils import CreateSpanOptions, SpanUtils
 from ...core.types import (
     CleanSpanData,
     Duration,
@@ -325,9 +326,31 @@ def _traced_execute(
 
         query_str = self._query_to_string(query, cursor)
 
-        # Create OpenTelemetry span
-        tracer = sdk.get_tracer()
-        span = tracer.start_span(
+        if sdk.mode == TuskDriftMode.REPLAY:
+            return handle_replay_mode(
+                replay_mode_handler=lambda: self._replay_execute(cursor, sdk, query_str, params),
+                no_op_request_handler=lambda: self._noop_execute(cursor),
+                is_server_request=False,
+            )
+
+        # RECORD mode
+        return handle_record_mode(
+            original_function_call=lambda: original_execute(query, params, **kwargs),
+            record_mode_handler=lambda is_pre_app_start: self._record_execute(
+                cursor, original_execute, sdk, query, query_str, params, is_pre_app_start, kwargs
+            ),
+            span_kind=OTelSpanKind.CLIENT,
+        )
+
+    def _noop_execute(self, cursor: Any) -> Any:
+        """Handle background requests in REPLAY mode - return cursor with empty mock data."""
+        cursor._mock_rows = []  # pyright: ignore
+        cursor._mock_index = 0  # pyright: ignore
+        return cursor
+
+    def _replay_execute(self, cursor: Any, sdk: TuskDrift, query_str: str, params: Any) -> Any:
+        """Handle REPLAY mode for execute - fetch mock from CLI."""
+        span_info = SpanUtils.create_span(CreateSpanOptions(
             name="psycopg.query",
             kind=OTelSpanKind.CLIENT,
             attributes={
@@ -338,80 +361,110 @@ def _traced_execute(
                 TdSpanAttributes.PACKAGE_TYPE: PackageType.PG.name,
                 TdSpanAttributes.IS_PRE_APP_START: not sdk.app_ready,
             },
-        )
+            is_pre_app_start=not sdk.app_ready,
+        ))
 
-        ctx = otel_context.get_current()
-        ctx_with_span = set_span_in_context(span, ctx)
-        token = otel_context.attach(ctx_with_span)
+        if not span_info:
+            raise RuntimeError("Error creating span in replay mode")
 
-        try:
-            span_context = span.get_span_context()
-            trace_id = format(span_context.trace_id, "032x")
-            span_id = format(span_context.span_id, "016x")
+        with SpanUtils.with_span(span_info):
+            mock_result = self._try_get_mock(
+                sdk, query_str, params,
+                span_info.trace_id, span_info.span_id, span_info.parent_span_id
+            )
 
-            parent_span = trace.get_current_span(ctx)
-            parent_span_id = None
-            if parent_span and parent_span.is_recording():
-                parent_ctx = parent_span.get_span_context()
-                parent_span_id = format(parent_ctx.span_id, "016x")
+            if mock_result is None:
+                is_pre_app_start = not sdk.app_ready
+                raise RuntimeError(
+                    f"[Tusk REPLAY] No mock found for psycopg execute query. "
+                    f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded during the trace capture. "
+                    f"Query: {query_str[:100]}..."
+                )
 
-            if sdk.mode == TuskDriftMode.REPLAY:
-                # Handle background requests (app ready + no parent span)
-                if sdk.app_ready and not parent_span_id:
-                    cursor._mock_rows = []  # pyright: ignore
-                    cursor._mock_index = 0  # pyright: ignore
-                    return cursor
-
-                mock_result = self._try_get_mock(sdk, query_str, params, trace_id, span_id, parent_span_id)
-
-                if mock_result is None:
-                    is_pre_app_start = not sdk.app_ready
-                    raise RuntimeError(
-                        f"[Tusk REPLAY] No mock found for psycopg execute query. "
-                        f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded during the trace capture. "
-                        f"Query: {query_str[:100]}..."
-                    )
+            self._mock_execute_with_data(cursor, mock_result)
+            span_info.span.end()
+            return cursor
 
-                self._mock_execute_with_data(cursor, mock_result)
-                return cursor
+    def _record_execute(
+        self,
+        cursor: Any,
+        original_execute: Any,
+        sdk: TuskDrift,
+        query: str,
+        query_str: str,
+        params: Any,
+        is_pre_app_start: bool,
+        kwargs: dict,
+    ) -> Any:
+        """Handle RECORD mode for execute - create span and execute query."""
+        span_info = SpanUtils.create_span(CreateSpanOptions(
+            name="psycopg.query",
+            kind=OTelSpanKind.CLIENT,
+            attributes={
+                TdSpanAttributes.NAME: "psycopg.query",
+                TdSpanAttributes.PACKAGE_NAME: "psycopg",
+                TdSpanAttributes.INSTRUMENTATION_NAME: "PsycopgInstrumentation",
+                TdSpanAttributes.SUBMODULE_NAME: "query",
+                TdSpanAttributes.PACKAGE_TYPE: PackageType.PG.name,
+                TdSpanAttributes.IS_PRE_APP_START: is_pre_app_start,
+            },
+            is_pre_app_start=is_pre_app_start,
+        ))
+
+        if not span_info:
+            # Fallback to original call if span creation fails
+            return original_execute(query, params, **kwargs)
 
-            # RECORD mode
-            time.time()
-            error = None
+        error = None
+        result = None
 
+        with SpanUtils.with_span(span_info):
             try:
                 result = original_execute(query, params, **kwargs)
                 return result
             except Exception as e:
                 error = e
                 raise
             finally:
-                if sdk.mode == TuskDriftMode.RECORD:
-                    self._finalize_query_span(
-                        span,
-                        cursor,
-                        query_str,
-                        params,
-                        error,
-                    )
-        finally:
-            otel_context.detach(token)
-            span.end()
+                self._finalize_query_span(
+                    span_info.span,
+                    cursor,
+                    query_str,
+                    params,
+                    error,
+                )
+                span_info.span.end()
 
     def _traced_executemany(
         self, cursor: Any, original_executemany: Any, sdk: TuskDrift, query: str, params_seq, **kwargs
     ) -> Any:
         """Traced cursor.executemany method."""
-        # Pass through if SDK is disabled
         if sdk.mode == TuskDriftMode.DISABLED:
             return original_executemany(query, params_seq, **kwargs)
 
-        # Convert query to string if it's a Composed SQL object
         query_str = self._query_to_string(query, cursor)
+        # Convert to list BEFORE executing to avoid iterator exhaustion
+        params_list = list(params_seq)
+
+        if sdk.mode == TuskDriftMode.REPLAY:
+            return handle_replay_mode(
+                replay_mode_handler=lambda: self._replay_executemany(cursor, sdk, query_str, params_list),
+                no_op_request_handler=lambda: self._noop_execute(cursor),
+                is_server_request=False,
+            )
+
+        # RECORD mode
+        return handle_record_mode(
+            original_function_call=lambda: original_executemany(query, params_list, **kwargs),
+            record_mode_handler=lambda is_pre_app_start: self._record_executemany(
+                cursor, original_executemany, sdk, query, query_str, params_list, is_pre_app_start, kwargs
+            ),
+            span_kind=OTelSpanKind.CLIENT,
+        )
 
-        # Create OpenTelemetry span
-        tracer = sdk.get_tracer()
-        span = tracer.start_span(
+    def _replay_executemany(self, cursor: Any, sdk: TuskDrift, query_str: str, params_list: list) -> Any:
+        """Handle REPLAY mode for executemany - fetch mock from CLI."""
+        span_info = SpanUtils.create_span(CreateSpanOptions(
             name="psycopg.query",
             kind=OTelSpanKind.CLIENT,
             attributes={
@@ -422,91 +475,82 @@ def _traced_executemany(
                 TdSpanAttributes.PACKAGE_TYPE: PackageType.PG.name,
                 TdSpanAttributes.IS_PRE_APP_START: not sdk.app_ready,
             },
-        )
+            is_pre_app_start=not sdk.app_ready,
+        ))
 
-        ctx = otel_context.get_current()
-        ctx_with_span = set_span_in_context(span, ctx)
-        token = otel_context.attach(ctx_with_span)
+        if not span_info:
+            raise RuntimeError("Error creating span in replay mode")
 
-        try:
-            span_context = span.get_span_context()
-            trace_id = format(span_context.trace_id, "032x")
-            span_id = format(span_context.span_id, "016x")
-
-            parent_span = trace.get_current_span(ctx)
-            parent_span_id = None
-            if parent_span and parent_span.is_recording():
-                parent_ctx = parent_span.get_span_context()
-                parent_span_id = format(parent_ctx.span_id, "016x")
-
-            # For executemany, we'll treat each parameter set as a batch
-            # REPLAY mode: Mock ALL queries (including pre-app-start)
-            if sdk.mode == TuskDriftMode.REPLAY:
-                # Handle background requests: App is ready + no parent span
-                # These are background jobs/health checks that run AFTER app startup
-                # They were never recorded, so return empty result
-                if sdk.app_ready and not parent_span_id:
-                    logger.debug("Background executemany request (app ready, no parent) - returning empty result")
-                    # Return the cursor
-                    cursor._mock_rows = []  # pyright: ignore
-                    cursor._mock_index = 0  # pyright: ignore
-                    return cursor
-
-                # For all other queries (pre-app-start OR within a request trace), get mock
-                # Convert params_seq to list for serialization
-                # Wrap in {"_batch": ...} to match the recording format
-                params_list = list(params_seq)
-                mock_result = self._try_get_mock(
-                    sdk,
-                    query_str,
-                    {"_batch": params_list},
-                    trace_id,
-                    span_id,
-                    parent_span_id,
+        with SpanUtils.with_span(span_info):
+            mock_result = self._try_get_mock(
+                sdk, query_str, {"_batch": params_list},
+                span_info.trace_id, span_info.span_id, span_info.parent_span_id
+            )
+
+            if mock_result is None:
+                is_pre_app_start = not sdk.app_ready
+                logger.error(
+                    f"No mock found for {'pre-app-start ' if is_pre_app_start else ''}psycopg executemany query in REPLAY mode: {query_str[:100]}"
+                )
+                raise RuntimeError(
+                    f"[Tusk REPLAY] No mock found for psycopg executemany query. "
+                    f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded during the trace capture. "
+                    f"Query: {query_str[:100]}..."
                 )
-                if mock_result is None:
-                    # In REPLAY mode, we MUST have a mock for ALL queries
-                    is_pre_app_start = not sdk.app_ready
-                    logger.error(
-                        f"No mock found for {'pre-app-start ' if is_pre_app_start else ''}psycopg executemany query in REPLAY mode: {query_str[:100]}"
-                    )
-                    raise RuntimeError(
-                        f"[Tusk REPLAY] No mock found for psycopg executemany query. "
-                        f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded during the trace capture. "
-                        f"Query: {query_str[:100]}..."
-                    )
 
-                # Mock execute by setting cursor internal state
-                self._mock_execute_with_data(cursor, mock_result)
-                return cursor  # psycopg3 executemany() returns cursor
+            self._mock_execute_with_data(cursor, mock_result)
+            span_info.span.end()
+            return cursor
 
-            # 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)
+    def _record_executemany(
+        self,
+        cursor: Any,
+        original_executemany: Any,
+        sdk: TuskDrift,
+        query: str,
+        query_str: str,
+        params_list: list,
+        is_pre_app_start: bool,
+        kwargs: dict,
+    ) -> Any:
+        """Handle RECORD mode for executemany - create span and execute query."""
+        span_info = SpanUtils.create_span(CreateSpanOptions(
+            name="psycopg.query",
+            kind=OTelSpanKind.CLIENT,
+            attributes={
+                TdSpanAttributes.NAME: "psycopg.query",
+                TdSpanAttributes.PACKAGE_NAME: "psycopg",
+                TdSpanAttributes.INSTRUMENTATION_NAME: "PsycopgInstrumentation",
+                TdSpanAttributes.SUBMODULE_NAME: "query",
+                TdSpanAttributes.PACKAGE_TYPE: PackageType.PG.name,
+                TdSpanAttributes.IS_PRE_APP_START: is_pre_app_start,
+            },
+            is_pre_app_start=is_pre_app_start,
+        ))
+
+        if not span_info:
+            # Fallback to original call if span creation fails
+            return original_executemany(query, params_list, **kwargs)
+
+        error = None
+        result = None
 
+        with SpanUtils.with_span(span_info):
             try:
                 result = original_executemany(query, params_list, **kwargs)
                 return result
             except Exception as e:
                 error = e
                 raise
             finally:
-                # 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:
-                    self._finalize_query_span(
-                        span,
-                        cursor,
-                        query_str,
-                        {"_batch": params_list},
-                        error,
-                    )
-        finally:
-            # Reset only span context (trace context is owned by parent)
-            otel_context.detach(token)
-            span.end()
+                self._finalize_query_span(
+                    span_info.span,
+                    cursor,
+                    query_str,
+                    {"_batch": params_list},
+                    error,
+                )
+                span_info.span.end()
 
     def _query_to_string(self, query: Any, cursor: Any) -> str:
         """Convert query to string."""