Allow removing extra FastAPI spans (#1258)

Author: alexmojaki

logfire/_internal/integrations/fastapi.py

@@ -4,6 +4,7 @@
 import inspect
 from collections.abc import Awaitable, Iterable
 from contextlib import AbstractContextManager, contextmanager
+from datetime import datetime, timezone
 from functools import lru_cache
 from typing import Any, Callable
 from weakref import WeakKeyDictionary
@@ -17,7 +18,8 @@
 from starlette.responses import Response
 from starlette.websockets import WebSocket
 
-from ..main import Logfire, set_user_attributes_on_raw_span
+from ..constants import ONE_SECOND_IN_NANOSECONDS
+from ..main import Logfire, NoopSpan, set_user_attributes_on_raw_span
 from ..stack_info import StackInfo, get_code_object_info
 from ..utils import handle_internal_errors, maybe_capture_server_headers
 
@@ -61,6 +63,7 @@ def instrument_fastapi(
     | None = None,
     excluded_urls: str | Iterable[str] | None = None,
     record_send_receive: bool = False,
+    extra_spans: bool = True,
     **opentelemetry_kwargs: Any,
 ) -> AbstractContextManager[None]:
     """Instrument a FastAPI app so that spans and logs are automatically created for each request.
@@ -96,6 +99,7 @@ def instrument_fastapi(
         registry[_app] = FastAPIInstrumentation(
             logfire_instance,
             request_attributes_mapper or _default_request_attributes_mapper,
+            extra_spans=extra_spans,
         )
 
     @contextmanager
@@ -158,16 +162,39 @@ def __init__(
             ],
             dict[str, Any] | None,
         ],
+        extra_spans: bool,
     ):
         self.logfire_instance = logfire_instance.with_settings(custom_scope_suffix='fastapi')
+        self.timestamp_generator = self.logfire_instance.config.advanced.ns_timestamp_generator
         self.request_attributes_mapper = request_attributes_mapper
+        self.extra_spans = extra_spans
+
+    @contextmanager
+    def pseudo_span(self, namespace: str, root_span: Span):
+        """Record start and end timestamps in the root span, and possibly exceptions."""
+
+        def set_timestamp(attribute_name: str):
+            dt = datetime.fromtimestamp(self.timestamp_generator() / ONE_SECOND_IN_NANOSECONDS, tz=timezone.utc)
+            value = dt.strftime('%Y-%m-%dT%H:%M:%S.%fZ')
+            root_span.set_attribute(f'fastapi.{namespace}.{attribute_name}', value)
+
+        set_timestamp('start_timestamp')
+        try:
+            try:
+                yield
+            finally:
+                # Record the end timestamp before recording exceptions.
+                set_timestamp('end_timestamp')
+        except Exception as exc:
+            root_span.record_exception(exc)
+            raise
 
     async def solve_dependencies(self, request: Request | WebSocket, original: Awaitable[Any]) -> Any:
         root_span = request.scope.get(LOGFIRE_SPAN_SCOPE_KEY)
         if not (root_span and root_span.is_recording()):
             return await original
 
-        with self.logfire_instance.span('FastAPI arguments') as span:
+        with self.logfire_instance.span('FastAPI arguments') if self.extra_spans else NoopSpan() as span:
             with handle_internal_errors:
                 if isinstance(request, Request):  # pragma: no branch
                     span.set_attribute('http.method', request.method)
@@ -180,7 +207,8 @@ async def solve_dependencies(self, request: Request | WebSocket, original: Await
                     set_user_attributes_on_raw_span(root_span, fastapi_route_attributes)
                     span.set_attributes(fastapi_route_attributes)
 
-            result: Any = await original
+            with self.pseudo_span('arguments', root_span):
+                result: Any = await original
 
             with handle_internal_errors:
                 solved_values: dict[str, Any]
@@ -228,7 +256,8 @@ def solved_with_new_values(new_values: dict[str, Any]) -> Any:
 
                 # request_attributes_mapper may have removed the errors, so we need .get() here.
                 if attributes.get('errors'):
-                    span.set_level('error')
+                    # Errors should imply a 422 response. 4xx errors are warnings, not errors.
+                    span.set_level('warn')
 
                 span.set_attributes(attributes)
                 for key in ('values', 'errors'):
@@ -246,20 +275,31 @@ async def run_endpoint_function(
         values: dict[str, Any],
         **kwargs: Any,
     ) -> Any:
-        callback = inspect.unwrap(dependant.call)
-        code = getattr(callback, '__code__', None)
-        stack_info: StackInfo = get_code_object_info(code) if code else {}
-        with self.logfire_instance.span(
-            '{method} {http.route} ({code.function})',
-            method=request.method,
-            # Using `http.route` prevents it from being scrubbed if it contains a word like 'secret'.
-            # We don't use `http.method` because some dashboards do things like count spans with
-            # both `http.method` and `http.route`.
-            **{'http.route': request.scope['route'].path},
-            **stack_info,
-            _level='debug',
-        ):
-            return await original_run_endpoint_function(dependant=dependant, values=values, **kwargs)
+        original = original_run_endpoint_function(dependant=dependant, values=values, **kwargs)
+        root_span = request.scope.get(LOGFIRE_SPAN_SCOPE_KEY)
+        if not (root_span and root_span.is_recording()):  # pragma: no cover
+            # This should never happen because we only get to this function after solve_dependencies
+            # passes the same check, just being paranoid.
+            return await original
+
+        if self.extra_spans:
+            callback = inspect.unwrap(dependant.call)
+            code = getattr(callback, '__code__', None)
+            stack_info: StackInfo = get_code_object_info(code) if code else {}
+            extra_span = self.logfire_instance.span(
+                '{method} {http.route} ({code.function})',
+                method=request.method,
+                # Using `http.route` prevents it from being scrubbed if it contains a word like 'secret'.
+                # We don't use `http.method` because some dashboards do things like count spans with
+                # both `http.method` and `http.route`.
+                **{'http.route': request.scope['route'].path},
+                **stack_info,
+                _level='debug',
+            )
+        else:
+            extra_span = NoopSpan()
+        with extra_span, self.pseudo_span('endpoint_function', root_span):
+            return await original
 
 
 def _default_request_attributes_mapper(

logfire/_internal/main.py

@@ -1035,6 +1035,7 @@ def instrument_fastapi(
         | None = None,
         excluded_urls: str | Iterable[str] | None = None,
         record_send_receive: bool = False,
+        extra_spans: bool = True,
         **opentelemetry_kwargs: Any,
     ) -> AbstractContextManager[None]:
         """Instrument a FastAPI app so that spans and logs are automatically created for each request.
@@ -1069,6 +1070,7 @@ def instrument_fastapi(
                 These are disabled by default to reduce overhead and the number of spans created,
                 since many can be created for a single request, and they are not often useful.
                 If enabled, they will be set to debug level, meaning they will usually still be hidden in the UI.
+            extra_spans: Whether to include the extra 'FastAPI arguments' and 'endpoint function' spans.
             opentelemetry_kwargs: Additional keyword arguments to pass to the OpenTelemetry FastAPI instrumentation.
 
         Returns:
@@ -1088,6 +1090,7 @@ def instrument_fastapi(
             request_attributes_mapper=request_attributes_mapper,
             excluded_urls=excluded_urls,
             record_send_receive=record_send_receive,
+            extra_spans=extra_spans,
             **opentelemetry_kwargs,
         )