fix formatting

Author: sohankshirsagar

drift/core/mode_utils.py

@@ -69,9 +69,7 @@ def handle_record_mode(
     # App is ready - check span context
     is_server_span = span_kind == OTelSpanKind.SERVER
 
-    if (not current_span_info and not is_server_span) or (
-        current_span_info and current_span_info.is_pre_app_start
-    ):
+    if (not current_span_info and not is_server_span) or (current_span_info and current_span_info.is_pre_app_start):
         # No span context and not a server request, OR within a pre-app-start span
         # Skip recording - call original function
         return original_function_call()

drift/instrumentation/django/middleware.py

@@ -75,15 +75,11 @@ def __call__(self, request: HttpRequest) -> HttpResponse:
         # RECORD mode - use handle_record_mode for consistent is_pre_app_start logic
         return handle_record_mode(
             original_function_call=lambda: self.get_response(request),
-            record_mode_handler=lambda is_pre_app_start: self._record_request(
-                request, sdk, is_pre_app_start
-            ),
+            record_mode_handler=lambda is_pre_app_start: self._record_request(request, sdk, is_pre_app_start),
             span_kind=OTelSpanKind.SERVER,
         )
 
-    def _handle_replay_request(
-        self, request: HttpRequest, sdk
-    ) -> HttpResponse:
+    def _handle_replay_request(self, request: HttpRequest, sdk) -> HttpResponse:
         """Handle request in REPLAY mode.
 
         Extracts trace ID from headers and sets up context for child spans.
@@ -98,23 +94,15 @@ def _handle_replay_request(
         """
         # Extract trace ID from headers (case-insensitive lookup)
         # Django stores headers in request.META
-        headers_lower = {
-            k.lower(): v for k, v in request.META.items() if k.startswith("HTTP_")
-        }
-        logger.info(
-            f"[DJANGO_MIDDLEWARE] REPLAY mode, headers: {list(headers_lower.keys())}"
-        )
+        headers_lower = {k.lower(): v for k, v in request.META.items() if k.startswith("HTTP_")}
+        logger.info(f"[DJANGO_MIDDLEWARE] REPLAY mode, headers: {list(headers_lower.keys())}")
         # Convert HTTP_X_TD_TRACE_ID -> x-td-trace-id
         replay_trace_id = headers_lower.get("http_x_td_trace_id")
-        logger.info(
-            f"[DJANGO_MIDDLEWARE] replay_trace_id from header: {replay_trace_id}"
-        )
+        logger.info(f"[DJANGO_MIDDLEWARE] replay_trace_id from header: {replay_trace_id}")
 
         if not replay_trace_id:
             # No trace context in REPLAY mode; proceed without span
-            logger.warning(
-                "[DJANGO_MIDDLEWARE] No replay_trace_id found in headers, proceeding without span"
-            )
+            logger.warning("[DJANGO_MIDDLEWARE] No replay_trace_id found in headers, proceeding without span")
             return self.get_response(request)
 
         # Set replay trace context
@@ -166,9 +154,7 @@ def _handle_replay_request(
             replay_trace_id_context.reset(replay_token)
             span_info.span.end()
 
-    def _record_request(
-        self, request: HttpRequest, sdk, is_pre_app_start: bool
-    ) -> HttpResponse:
+    def _record_request(self, request: HttpRequest, sdk, is_pre_app_start: bool) -> HttpResponse:
         """Handle request in RECORD mode.
 
         Creates a span, processes the request, and captures the span.
@@ -188,9 +174,7 @@ def _record_request(
 
             sampling_rate = sdk.get_sampling_rate()
             if not should_sample(sampling_rate, is_app_ready=True):
-                logger.debug(
-                    f"[Django] Request not sampled (rate={sampling_rate}), path={request.path}"
-                )
+                logger.debug(f"[Django] Request not sampled (rate={sampling_rate}), path={request.path}")
                 return self.get_response(request)
 
         start_time_ns = time.time_ns()
@@ -240,9 +224,7 @@ def _record_request(
 
         request_headers = extract_headers(request.META)
 
-        if self.transform_engine and self.transform_engine.should_drop_inbound_request(
-            method, target, request_headers
-        ):
+        if self.transform_engine and self.transform_engine.should_drop_inbound_request(method, target, request_headers):
             # Reset context before early return
             span_kind_context.reset(span_kind_token)
             span_info.span.end()
@@ -289,9 +271,7 @@ def process_view(
             if route:
                 request._drift_route_template = route  # type: ignore
 
-    def _capture_span(
-        self, request: HttpRequest, response: HttpResponse, span_info: SpanInfo
-    ) -> None:
+    def _capture_span(self, request: HttpRequest, response: HttpResponse, span_info: SpanInfo) -> None:
         """Create and collect a span from request/response data.
 
         Args:
@@ -458,9 +438,7 @@ def dict_to_schema_merges(merges_dict):
 
             sdk.collect_span(clean_span)
 
-    def _capture_error_span(
-        self, request: HttpRequest, exception: Exception, span_info: SpanInfo
-    ) -> None:
+    def _capture_error_span(self, request: HttpRequest, exception: Exception, span_info: SpanInfo) -> None:
         """Create and collect an error span.
 
         Args:

drift/instrumentation/fastapi/instrumentation.py

@@ -197,9 +197,7 @@ async def wrapped_send(message: dict[str, Any]) -> None:
             response_data["status_message"] = _get_status_message(message.get("status", 200))
             raw_headers = message.get("headers", [])
             response_data["headers"] = {
-                k.decode("utf-8", errors="replace") if isinstance(k, bytes) else k: v.decode(
-                    "utf-8", errors="replace"
-                )
+                k.decode("utf-8", errors="replace") if isinstance(k, bytes) else k: v.decode("utf-8", errors="replace")
                 if isinstance(v, bytes)
                 else v
                 for k, v in raw_headers
@@ -324,9 +322,7 @@ async def wrapped_send(message: dict[str, Any]) -> None:
             response_data["status_message"] = _get_status_message(message.get("status", 200))
             raw_headers = message.get("headers", [])
             response_data["headers"] = {
-                k.decode("utf-8", errors="replace") if isinstance(k, bytes) else k: v.decode(
-                    "utf-8", errors="replace"
-                )
+                k.decode("utf-8", errors="replace") if isinstance(k, bytes) else k: v.decode("utf-8", errors="replace")
                 if isinstance(v, bytes)
                 else v
                 for k, v in raw_headers

drift/instrumentation/psycopg/instrumentation.py

@@ -350,27 +350,28 @@ def _noop_execute(self, cursor: Any) -> Any:
 
     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={
-                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: not sdk.app_ready,
-            },
-            is_pre_app_start=not sdk.app_ready,
-        ))
+        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: not sdk.app_ready,
+                },
+                is_pre_app_start=not sdk.app_ready,
+            )
+        )
 
         if not span_info:
             raise RuntimeError("Error creating span in replay mode")
 
         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
+                sdk, query_str, params, span_info.trace_id, span_info.span_id, span_info.parent_span_id
             )
 
             if mock_result is None:
@@ -397,19 +398,21 @@ def _record_execute(
         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,
-        ))
+        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
@@ -464,27 +467,28 @@ def _traced_executemany(
 
     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={
-                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: not sdk.app_ready,
-            },
-            is_pre_app_start=not sdk.app_ready,
-        ))
+        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: not sdk.app_ready,
+                },
+                is_pre_app_start=not sdk.app_ready,
+            )
+        )
 
         if not span_info:
             raise RuntimeError("Error creating span in replay mode")
 
         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
+                sdk, query_str, {"_batch": params_list}, span_info.trace_id, span_info.span_id, span_info.parent_span_id
             )
 
             if mock_result is None:
@@ -514,19 +518,21 @@ def _record_executemany(
         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,
-        ))
+        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