Skip to content

Commit

Permalink
Use OPEN_SPANS for all spans
Browse files Browse the repository at this point in the history
  • Loading branch information
dmontagu committed Dec 18, 2024
1 parent 539311e commit 6f48eea
Show file tree
Hide file tree
Showing 7 changed files with 110 additions and 103 deletions.
6 changes: 2 additions & 4 deletions logfire-api/logfire_api/_internal/config.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,11 @@ from .exporters.quiet_metrics import QuietMetricExporter as QuietMetricExporter
from .exporters.remove_pending import RemovePendingSpansExporter as RemovePendingSpansExporter
from .exporters.test import TestExporter as TestExporter
from .integrations.executors import instrument_executors as instrument_executors
from .main import FastLogfireSpan as FastLogfireSpan, Logfire as Logfire, LogfireSpan as LogfireSpan
from .main import Logfire as Logfire
from .metrics import ProxyMeterProvider as ProxyMeterProvider
from .scrubbing import BaseScrubber as BaseScrubber, NOOP_SCRUBBER as NOOP_SCRUBBER, Scrubber as Scrubber, ScrubbingOptions as ScrubbingOptions
from .stack_info import warn_at_user_stacklevel as warn_at_user_stacklevel
from .tracer import PendingSpanProcessor as PendingSpanProcessor, ProxyTracerProvider as ProxyTracerProvider
from .tracer import OPEN_SPANS as OPEN_SPANS, PendingSpanProcessor as PendingSpanProcessor, ProxyTracerProvider as ProxyTracerProvider
from .utils import SeededRandomIdGenerator as SeededRandomIdGenerator, UnexpectedResponse as UnexpectedResponse, ensure_data_dir_exists as ensure_data_dir_exists, handle_internal_errors as handle_internal_errors, read_toml_file as read_toml_file, suppress_instrumentation as suppress_instrumentation
from _typeshed import Incomplete
from dataclasses import dataclass
Expand All @@ -30,9 +30,7 @@ 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

OPEN_SPANS: WeakSet[LogfireSpan | FastLogfireSpan]
CREDENTIALS_FILENAME: str
COMMON_REQUEST_HEADERS: Incomplete
PROJECT_NAME_PATTERN: str
Expand Down
10 changes: 5 additions & 5 deletions logfire-api/logfire_api/_internal/main.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@ import requests
from . import async_ as async_
from ..version import VERSION as VERSION
from .auto_trace import AutoTraceModule as AutoTraceModule, install_auto_tracing as install_auto_tracing
from .config import GLOBAL_CONFIG as GLOBAL_CONFIG, LogfireConfig as LogfireConfig, OPEN_SPANS as OPEN_SPANS
from .config import GLOBAL_CONFIG as GLOBAL_CONFIG, LogfireConfig as LogfireConfig
from .config_params import PydanticPluginRecordValues as PydanticPluginRecordValues
from .constants import ATTRIBUTES_JSON_SCHEMA_KEY as ATTRIBUTES_JSON_SCHEMA_KEY, ATTRIBUTES_LOG_LEVEL_NUM_KEY as ATTRIBUTES_LOG_LEVEL_NUM_KEY, ATTRIBUTES_MESSAGE_KEY as ATTRIBUTES_MESSAGE_KEY, ATTRIBUTES_MESSAGE_TEMPLATE_KEY as ATTRIBUTES_MESSAGE_TEMPLATE_KEY, ATTRIBUTES_SAMPLE_RATE_KEY as ATTRIBUTES_SAMPLE_RATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY as ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_TAGS_KEY as ATTRIBUTES_TAGS_KEY, ATTRIBUTES_VALIDATION_ERROR_KEY as ATTRIBUTES_VALIDATION_ERROR_KEY, DISABLE_CONSOLE_KEY as DISABLE_CONSOLE_KEY, LEVEL_NUMBERS as LEVEL_NUMBERS, LevelName as LevelName, NULL_ARGS_KEY as NULL_ARGS_KEY, OTLP_MAX_INT_SIZE as OTLP_MAX_INT_SIZE, log_level_attributes as log_level_attributes
from .constants import ATTRIBUTES_JSON_SCHEMA_KEY as ATTRIBUTES_JSON_SCHEMA_KEY, ATTRIBUTES_LOG_LEVEL_NUM_KEY as ATTRIBUTES_LOG_LEVEL_NUM_KEY, ATTRIBUTES_MESSAGE_KEY as ATTRIBUTES_MESSAGE_KEY, ATTRIBUTES_MESSAGE_TEMPLATE_KEY as ATTRIBUTES_MESSAGE_TEMPLATE_KEY, ATTRIBUTES_SAMPLE_RATE_KEY as ATTRIBUTES_SAMPLE_RATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY as ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_TAGS_KEY as ATTRIBUTES_TAGS_KEY, DISABLE_CONSOLE_KEY as DISABLE_CONSOLE_KEY, LEVEL_NUMBERS as LEVEL_NUMBERS, LevelName as LevelName, NULL_ARGS_KEY as NULL_ARGS_KEY, OTLP_MAX_INT_SIZE as OTLP_MAX_INT_SIZE, log_level_attributes as log_level_attributes
from .formatter import logfire_format as logfire_format, logfire_format_with_magic as logfire_format_with_magic
from .instrument import instrument as instrument
from .integrations.asgi import ASGIApp as ASGIApp, ASGIInstrumentKwargs as ASGIInstrumentKwargs
Expand All @@ -29,7 +29,7 @@ from .json_encoder import logfire_json_dumps as logfire_json_dumps
from .json_schema import JsonSchemaProperties as JsonSchemaProperties, attributes_json_schema as attributes_json_schema, attributes_json_schema_properties as attributes_json_schema_properties, create_json_schema as create_json_schema
from .metrics import ProxyMeterProvider as ProxyMeterProvider
from .stack_info import get_user_stack_info as get_user_stack_info
from .tracer import ProxyTracerProvider as ProxyTracerProvider
from .tracer import ProxyTracerProvider as ProxyTracerProvider, record_exception as record_exception, set_exception_status as set_exception_status
from .utils import SysExcInfo as SysExcInfo, get_version as get_version, handle_internal_errors as handle_internal_errors, log_internal_error as log_internal_error, uniquify_sequence as uniquify_sequence
from django.http import HttpRequest as HttpRequest, HttpResponse as HttpResponse
from fastapi import FastAPI
Expand Down Expand Up @@ -1034,7 +1034,7 @@ class LogfireSpan(ReadableSpan):
def message(self) -> str: ...
@message.setter
def message(self, message: str): ...
def end(self) -> None:
def end(self, end_time: int | None = None) -> None:
"""Sets the current time as the span's end time.
The span's end time is the wall time at which the operation finished.
Expand Down Expand Up @@ -1093,7 +1093,7 @@ class NoopSpan:
def is_recording(self) -> bool: ...
AttributesValueType = TypeVar('AttributesValueType', bound=Any | otel_types.AttributeValue)

def user_attributes(attributes: dict[str, Any]) -> dict[str, otel_types.AttributeValue]:
def prepare_otlp_attributes(attributes: dict[str, Any]) -> dict[str, otel_types.AttributeValue]:
"""Prepare attributes for sending to OpenTelemetry.
This will convert any non-OpenTelemetry compatible types to JSON.
Expand Down
14 changes: 11 additions & 3 deletions logfire-api/logfire_api/_internal/tracer.pyi
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import opentelemetry.trace as trace_api
from .config import LogfireConfig as LogfireConfig
from .constants import ATTRIBUTES_MESSAGE_KEY as ATTRIBUTES_MESSAGE_KEY, ATTRIBUTES_PENDING_SPAN_REAL_PARENT_KEY as ATTRIBUTES_PENDING_SPAN_REAL_PARENT_KEY, ATTRIBUTES_SAMPLE_RATE_KEY as ATTRIBUTES_SAMPLE_RATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY as ATTRIBUTES_SPAN_TYPE_KEY, PENDING_SPAN_NAME_SUFFIX as PENDING_SPAN_NAME_SUFFIX
from .constants import ATTRIBUTES_MESSAGE_KEY as ATTRIBUTES_MESSAGE_KEY, ATTRIBUTES_PENDING_SPAN_REAL_PARENT_KEY as ATTRIBUTES_PENDING_SPAN_REAL_PARENT_KEY, ATTRIBUTES_SAMPLE_RATE_KEY as ATTRIBUTES_SAMPLE_RATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY as ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_VALIDATION_ERROR_KEY as ATTRIBUTES_VALIDATION_ERROR_KEY, PENDING_SPAN_NAME_SUFFIX as PENDING_SPAN_NAME_SUFFIX, log_level_attributes as log_level_attributes
from .utils import handle_internal_errors as handle_internal_errors
from _typeshed import Incomplete
from dataclasses import dataclass
from opentelemetry import context as context_api
Expand All @@ -13,7 +14,9 @@ from opentelemetry.trace.status import Status, StatusCode
from opentelemetry.util import types as otel_types
from threading import Lock
from typing import Any, Callable, Mapping, Sequence
from weakref import WeakKeyDictionary
from weakref import WeakKeyDictionary, WeakSet

OPEN_SPANS: WeakSet[_LogfireWrappedSpan]

@dataclass
class ProxyTracerProvider(TracerProvider):
Expand All @@ -33,10 +36,11 @@ class ProxyTracerProvider(TracerProvider):
def force_flush(self, timeout_millis: int = 30000) -> bool: ...

@dataclass
class _MaybeDeterministicTimestampSpan(trace_api.Span, ReadableSpan):
class _LogfireWrappedSpan(trace_api.Span, ReadableSpan):
"""Span that overrides end() to use a timestamp generator if one was provided."""
span: Span
ns_timestamp_generator: Callable[[], int]
def __post_init__(self) -> None: ...
def end(self, end_time: int | None = None) -> None: ...
def get_span_context(self) -> SpanContext: ...
def set_attributes(self, attributes: dict[str, otel_types.AttributeValue]) -> None: ...
Expand All @@ -47,6 +51,7 @@ class _MaybeDeterministicTimestampSpan(trace_api.Span, ReadableSpan):
def is_recording(self) -> bool: ...
def set_status(self, status: Status | StatusCode, description: str | None = None) -> None: ...
def record_exception(self, exception: BaseException, attributes: otel_types.Attributes = None, timestamp: int | None = None, escaped: bool = False) -> None: ...
__hash__ = ...
def __getattr__(self, name: str) -> Any: ...

@dataclass
Expand Down Expand Up @@ -85,3 +90,6 @@ def should_sample(span_context: SpanContext, attributes: Mapping[str, otel_types
This is used to sample spans that are not sampled by the OTEL sampler.
"""
def get_sample_rate_from_attributes(attributes: otel_types.Attributes) -> float | None: ...
def record_exception(span: trace_api.Span, exception: BaseException, *, attributes: otel_types.Attributes = None, timestamp: int | None = None, escaped: bool = False) -> None:
"""Similar to the OTEL SDK Span.record_exception method, with our own additions."""
def set_exception_status(span: trace_api.Span, exception: BaseException): ...
6 changes: 3 additions & 3 deletions logfire/_internal/backfill.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
)
from .exporters.file import FileSpanExporter
from .formatter import logfire_format
from .main import user_attributes
from .main import prepare_otlp_attributes
from .scrubbing import Scrubber

try:
Expand Down Expand Up @@ -150,7 +150,7 @@ def write(self, data: Union[Log, Span]) -> None:
)
else:
parent_context = None # pragma: no cover
otlp_attributes = user_attributes(data.attributes)
otlp_attributes = prepare_otlp_attributes(data.attributes)

if data.formatted_msg is None: # pragma: no cover
formatted_message = logfire_format(data.msg_template, data.attributes, self.scrubber)
Expand Down Expand Up @@ -196,7 +196,7 @@ def write(self, data: Union[Log, Span]) -> None:
start_timestamp = data.start_timestamp
if start_timestamp.tzinfo is None: # pragma: no branch
start_timestamp = start_timestamp.replace(tzinfo=timezone.utc)
otlp_attributes = user_attributes(data.log_attributes)
otlp_attributes = prepare_otlp_attributes(data.log_attributes)
if data.formatted_msg is None: # pragma: no branch
formatted_message = logfire_format(data.msg_template, data.log_attributes, self.scrubber)
else: # pragma: no cover
Expand Down
7 changes: 2 additions & 5 deletions logfire/_internal/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence, TypedDict, cast
from urllib.parse import urljoin
from uuid import uuid4
from weakref import WeakSet

import requests
from opentelemetry import trace
Expand Down Expand Up @@ -87,7 +86,7 @@
from .metrics import ProxyMeterProvider
from .scrubbing import NOOP_SCRUBBER, BaseScrubber, Scrubber, ScrubbingOptions
from .stack_info import warn_at_user_stacklevel
from .tracer import PendingSpanProcessor, ProxyTracerProvider
from .tracer import OPEN_SPANS, PendingSpanProcessor, ProxyTracerProvider
from .utils import (
SeededRandomIdGenerator,
UnexpectedResponse,
Expand All @@ -98,10 +97,8 @@
)

if TYPE_CHECKING:
from .main import FastLogfireSpan, Logfire, LogfireSpan
from .main import Logfire

# NOTE: this WeakSet is the reason that FastLogfireSpan.__slots__ has a __weakref__ slot.
OPEN_SPANS: WeakSet[LogfireSpan | FastLogfireSpan] = WeakSet()

CREDENTIALS_FILENAME = 'logfire_credentials.json'
"""Default base URL for the Logfire API."""
Expand Down
Loading

0 comments on commit 6f48eea

Please sign in to comment.