Add APIs for getting the current logfire span

logfire-api/logfire_api/__init__.py

@@ -18,6 +18,12 @@
 
         def configure(*args, **kwargs): ...
 
+        def current_span(*args, **kwargs):
+            return MagicMock()
+
+        def current_logfire_span(*args, **kwargs):
+            return MagicMock()
+
         class LogfireSpan:
             def __getattr__(self, attr):
                 return MagicMock()  # pragma: no cover

logfire-api/logfire_api/__init__.pyi

@@ -6,13 +6,14 @@ from ._internal.exporters.file import load_file as load_spans_from_file
 from ._internal.main import Logfire as Logfire, LogfireSpan as LogfireSpan
 from ._internal.scrubbing import ScrubMatch as ScrubMatch, ScrubbingOptions as ScrubbingOptions
 from ._internal.utils import suppress_instrumentation as suppress_instrumentation
+from .current_span import current_logfire_span, current_span
 from .integrations.logging import LogfireLoggingHandler as LogfireLoggingHandler
 from .integrations.structlog import LogfireProcessor as StructlogProcessor
 from .version import VERSION as VERSION
 from logfire.sampling import SamplingOptions as SamplingOptions
 from typing import Any
 
-__all__ = ['Logfire', 'LogfireSpan', 'LevelName', 'AdvancedOptions', 'ConsoleOptions', 'CodeSource', 'PydanticPlugin', 'configure', 'span', 'instrument', 'log', 'trace', 'debug', 'notice', 'info', 'warn', 'error', 'exception', 'fatal', 'force_flush', 'log_slow_async_callbacks', 'install_auto_tracing', 'instrument_asgi', 'instrument_wsgi', 'instrument_pydantic', 'instrument_fastapi', 'instrument_openai', 'instrument_anthropic', 'instrument_asyncpg', 'instrument_httpx', 'instrument_celery', 'instrument_requests', 'instrument_psycopg', 'instrument_django', 'instrument_flask', 'instrument_starlette', 'instrument_aiohttp_client', 'instrument_sqlalchemy', 'instrument_sqlite3', 'instrument_aws_lambda', 'instrument_redis', 'instrument_pymongo', 'instrument_mysql', 'instrument_system_metrics', 'AutoTraceModule', 'with_tags', 'with_settings', 'suppress_scopes', 'shutdown', 'load_spans_from_file', 'no_auto_trace', 'ScrubMatch', 'ScrubbingOptions', 'VERSION', 'suppress_instrumentation', 'StructlogProcessor', 'LogfireLoggingHandler', 'loguru_handler', 'SamplingOptions', 'MetricsOptions']
+__all__ = ['Logfire', 'LogfireSpan', 'LevelName', 'AdvancedOptions', 'ConsoleOptions', 'CodeSource', 'PydanticPlugin', 'current_span', 'current_logfire_span', 'configure', 'span', 'instrument', 'log', 'trace', 'debug', 'notice', 'info', 'warn', 'error', 'exception', 'fatal', 'force_flush', 'log_slow_async_callbacks', 'install_auto_tracing', 'instrument_asgi', 'instrument_wsgi', 'instrument_pydantic', 'instrument_fastapi', 'instrument_openai', 'instrument_anthropic', 'instrument_asyncpg', 'instrument_httpx', 'instrument_celery', 'instrument_requests', 'instrument_psycopg', 'instrument_django', 'instrument_flask', 'instrument_starlette', 'instrument_aiohttp_client', 'instrument_sqlalchemy', 'instrument_sqlite3', 'instrument_aws_lambda', 'instrument_redis', 'instrument_pymongo', 'instrument_mysql', 'instrument_system_metrics', 'AutoTraceModule', 'with_tags', 'with_settings', 'suppress_scopes', 'shutdown', 'load_spans_from_file', 'no_auto_trace', 'ScrubMatch', 'ScrubbingOptions', 'VERSION', 'suppress_instrumentation', 'StructlogProcessor', 'LogfireLoggingHandler', 'loguru_handler', 'SamplingOptions', 'MetricsOptions']
 
 DEFAULT_LOGFIRE_INSTANCE = Logfire()
 span = DEFAULT_LOGFIRE_INSTANCE.span

logfire-api/logfire_api/_internal/config.pyi

@@ -30,9 +30,10 @@ from opentelemetry.sdk.trace.id_generator import IdGenerator
 from pathlib import Path
 from typing import Any, Callable, Literal, Sequence, TypedDict
 from typing_extensions import Self, Unpack
-from weakref import WeakSet
+from weakref import WeakSet, WeakValueDictionary
 
 OPEN_SPANS: WeakSet[LogfireSpan | FastLogfireSpan]
+OPEN_SPANS_BY_ID: WeakValueDictionary[tuple[int, int], LogfireSpan | FastLogfireSpan]
 CREDENTIALS_FILENAME: str
 COMMON_REQUEST_HEADERS: Incomplete
 PROJECT_NAME_PATTERN: str

logfire-api/logfire_api/current_span.pyi

@@ -0,0 +1,23 @@
+from __future__ import annotations as _annotations
+
+from opentelemetry.trace import get_current_span
+
+from logfire import LogfireSpan
+
+__all__ = ('current_span', 'current_logfire_span')
+
+
+current_span = get_current_span
+
+
+def current_logfire_span() -> LogfireSpan:
+    """Return the LogfireSpan corresponding to the current otel span.
+
+    If the current otel span was not created as a LogfireSpan, we warn and return
+    something API-compatible which delegates to the otel span as much as possible.
+
+    Note: If we eventually rework the SDK so `opentelemetry.trace.get_current_span` returns a `LogfireSpan`, we should
+    make this an alias for `current_span` and deprecate this method. There are some good reasons to do that, but there
+    are also some good reasons not to, such as reducing overhead in calls made by third-party instrumentations.
+    """
+    ...

logfire/__init__.py

@@ -21,6 +21,7 @@
 from ._internal.main import Logfire, LogfireSpan
 from ._internal.scrubbing import ScrubbingOptions, ScrubMatch
 from ._internal.utils import suppress_instrumentation
+from .current_span import current_logfire_span, current_span
 from .integrations.logging import LogfireLoggingHandler
 from .integrations.structlog import LogfireProcessor as StructlogProcessor
 from .version import VERSION
@@ -101,6 +102,8 @@ def loguru_handler() -> Any:
     'ConsoleOptions',
     'CodeSource',
     'PydanticPlugin',
+    'current_span',
+    'current_logfire_span',
     'configure',
     'span',
     'instrument',

logfire/_internal/config.py

@@ -16,7 +16,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence, TypedDict, cast
 from urllib.parse import urljoin
 from uuid import uuid4
-from weakref import WeakSet
+from weakref import WeakSet, WeakValueDictionary
 
 import requests
 from opentelemetry import trace
@@ -102,6 +102,7 @@
 
 # NOTE: this WeakSet is the reason that FastLogfireSpan.__slots__ has a __weakref__ slot.
 OPEN_SPANS: WeakSet[LogfireSpan | FastLogfireSpan] = WeakSet()
+OPEN_LOGFIRE_SPANS_BY_ID: WeakValueDictionary[tuple[int, int], LogfireSpan] = WeakValueDictionary()
 
 CREDENTIALS_FILENAME = 'logfire_credentials.json'
 """Default base URL for the Logfire API."""

logfire/_internal/main.py

@@ -34,7 +34,7 @@
 from ..version import VERSION
 from . import async_
 from .auto_trace import AutoTraceModule, install_auto_tracing
-from .config import GLOBAL_CONFIG, OPEN_SPANS, LogfireConfig
+from .config import GLOBAL_CONFIG, OPEN_LOGFIRE_SPANS_BY_ID, OPEN_SPANS, LogfireConfig
 from .config_params import PydanticPluginRecordValues
 from .constants import (
     ATTRIBUTES_JSON_SCHEMA_KEY,
@@ -1908,6 +1908,8 @@ def __enter__(self) -> LogfireSpan:
             if self._token is None:  # pragma: no branch
                 self._token = context_api.attach(trace_api.set_span_in_context(self._span))
 
+            span_context = self._span.get_span_context()
+            OPEN_LOGFIRE_SPANS_BY_ID[(span_context.trace_id, span_context.span_id)] = self
             OPEN_SPANS.add(self)
 
         return self
@@ -1917,6 +1919,9 @@ def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseExceptio
         if self._token is None:  # pragma: no cover
             return
 
+        if self._span is not None:
+            span_context = self._span.get_span_context()
+            OPEN_LOGFIRE_SPANS_BY_ID.pop((span_context.trace_id, span_context.span_id), None)
         OPEN_SPANS.remove(self)
 
         context_api.detach(self._token)

logfire/current_span.py

@@ -0,0 +1,51 @@
+from __future__ import annotations as _annotations
+
+import warnings
+
+from opentelemetry.trace import get_current_span
+from opentelemetry.trace.span import Span
+
+from logfire._internal.config import OPEN_LOGFIRE_SPANS_BY_ID
+from logfire._internal.main import LogfireSpan, NoopSpan
+
+__all__ = ('current_span', 'current_logfire_span')
+
+
+class _BestEffortSpan:
+    def __init__(self, span: Span):
+        self.__span = span
+        self.__noop_span = NoopSpan()
+
+    def __getattr__(self, name: str):
+        try:
+            return getattr(self.__span, name)
+        except AttributeError:
+            value = getattr(self.__noop_span, name)
+            # Emit the warning _after_ grabbing the value so we don't emit a warning if an AttributeError will be raised
+            warnings.warn(
+                'A logfire-specific attribute is being accessed on a non-logfire span,'
+                ' the value is not meaningful and method calls will not do anything.',
+                stacklevel=2,
+            )
+            return value
+
+
+current_span = get_current_span
+
+
+def current_logfire_span() -> LogfireSpan:
+    """Return the LogfireSpan corresponding to the current otel span.
+
+    If the current otel span was not created as a LogfireSpan, we warn and return
+    something API-compatible which delegates to the otel span as much as possible.
+
+    Note: If we eventually rework the SDK so `opentelemetry.trace.get_current_span` returns a `LogfireSpan`, we should
+    make this an alias for `current_span` and deprecate this method. There are some good reasons to do that, but there
+    are also some good reasons not to, such as reducing overhead in calls made by third-party instrumentations.
+    """
+    otel_span = get_current_span()
+    span_context = otel_span.get_span_context()
+    logfire_span = OPEN_LOGFIRE_SPANS_BY_ID.get((span_context.trace_id, span_context.span_id))
+    if isinstance(logfire_span, LogfireSpan):
+        return logfire_span
+    return _BestEffortSpan(otel_span)  # type: ignore

tests/test_logfire_api.py

@@ -111,6 +111,14 @@ def test_runtime(logfire_api_factory: Callable[[], ModuleType], module_name: str
     logfire_api.PydanticPlugin()
     logfire__all__.remove('PydanticPlugin')
 
+    assert hasattr(logfire_api, 'current_span')
+    logfire_api.current_span()
+    logfire__all__.remove('current_span')
+
+    assert hasattr(logfire_api, 'current_logfire_span')
+    logfire_api.current_logfire_span()
+    logfire__all__.remove('current_logfire_span')
+
     assert hasattr(logfire_api, 'ScrubMatch')
     logfire_api.ScrubMatch(path='test', value='test', pattern_match='test')
     logfire__all__.remove('ScrubMatch')