From 900147559d4cb69c22b2799dbd66f5e4ae5bb3a0 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 14 Nov 2024 15:44:31 +0100 Subject: [PATCH] Add span links (#587) Co-authored-by: Alex Hall --- logfire/_internal/exporters/test.py | 28 +++---- logfire/_internal/main.py | 16 +++- logfire/_internal/tracer.py | 3 + tests/test_logfire.py | 113 ++++++++++++++++++++++++++++ 4 files changed, 143 insertions(+), 17 deletions(-) diff --git a/logfire/_internal/exporters/test.py b/logfire/_internal/exporters/test.py index 7b0498bb7..590b91401 100644 --- a/logfire/_internal/exporters/test.py +++ b/logfire/_internal/exporters/test.py @@ -88,11 +88,15 @@ def build_attributes(attributes: Mapping[str, Any] | None) -> dict[str, Any] | N attributes['telemetry.sdk.version'] = '0.0.0' return attributes + def build_context(context: trace.SpanContext) -> dict[str, Any]: + return {'trace_id': context.trace_id, 'span_id': context.span_id, 'is_remote': context.is_remote} + + def build_link(link: trace.Link) -> dict[str, Any]: + context = link.context or trace.INVALID_SPAN_CONTEXT + return {'context': build_context(context), 'attributes': build_attributes(link.attributes)} + def build_event(event: Event) -> dict[str, Any]: - res: dict[str, Any] = { - 'name': event.name, - 'timestamp': event.timestamp, - } + res: dict[str, Any] = {'name': event.name, 'timestamp': event.timestamp} if event.attributes: # pragma: no branch res['attributes'] = attributes = dict(event.attributes) if SpanAttributes.EXCEPTION_STACKTRACE in attributes: @@ -116,23 +120,15 @@ def build_span(span: ReadableSpan) -> dict[str, Any]: context = span.context or trace.INVALID_SPAN_CONTEXT res: dict[str, Any] = { 'name': span.name, - 'context': { - 'trace_id': context.trace_id, - 'span_id': context.span_id, - 'is_remote': context.is_remote, - }, - 'parent': { - 'trace_id': span.parent.trace_id, - 'span_id': span.parent.span_id, - 'is_remote': span.parent.is_remote, - } - if span.parent - else None, + 'context': build_context(context), + 'parent': build_context(span.parent) if span.parent else None, 'start_time': span.start_time, 'end_time': span.end_time, **build_instrumentation_scope(span), 'attributes': build_attributes(span.attributes), } + if span.links: + res['links'] = [build_link(link) for link in span.links] if span.events: res['events'] = [build_event(event) for event in span.events] if include_resources: diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 20d5d4686..4d2e1eaa5 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -15,7 +15,7 @@ from opentelemetry.metrics import CallbackT, Counter, Histogram, UpDownCounter from opentelemetry.sdk.trace import ReadableSpan, Span from opentelemetry.semconv.trace import SpanAttributes -from opentelemetry.trace import StatusCode, Tracer +from opentelemetry.trace import SpanContext, StatusCode, Tracer from opentelemetry.util import types as otel_types from typing_extensions import LiteralString, ParamSpec @@ -154,6 +154,7 @@ def _span( _tags: Sequence[str] | None = None, _span_name: str | None = None, _level: LevelName | int | None = None, + _links: Sequence[tuple[SpanContext, otel_types.Attributes]] = (), ) -> LogfireSpan: try: stack_info = get_user_stack_info() @@ -200,6 +201,7 @@ def _span( otlp_attributes, self._spans_tracer, json_schema_properties, + links=_links, ) except Exception: log_internal_error() @@ -492,6 +494,7 @@ def span( _tags: Sequence[str] | None = None, _span_name: str | None = None, _level: LevelName | None = None, + _links: Sequence[tuple[SpanContext, otel_types.Attributes]] = (), **attributes: Any, ) -> LogfireSpan: """Context manager for creating a span. @@ -510,6 +513,7 @@ def span( _span_name: The span name. If not provided, the `msg_template` will be used. _tags: An optional sequence of tags to include in the span. _level: An optional log level name. + _links: An optional sequence of links to other spans. Each link is a tuple of a span context and attributes. attributes: The arguments to include in the span and format the message template with. Attributes starting with an underscore are not allowed. """ @@ -521,6 +525,7 @@ def span( _tags=_tags, _span_name=_span_name, _level=_level, + _links=_links, ) def instrument( @@ -1752,11 +1757,13 @@ def __init__( otlp_attributes: dict[str, otel_types.AttributeValue], tracer: Tracer, json_schema_properties: JsonSchemaProperties, + links: Sequence[tuple[SpanContext, otel_types.Attributes]], ) -> None: self._span_name = span_name self._otlp_attributes = otlp_attributes self._tracer = tracer self._json_schema_properties = json_schema_properties + self._links = list(trace_api.Link(context=context, attributes=attributes) for context, attributes in links) self._added_attributes = False self._end_on_exit: bool | None = None @@ -1776,6 +1783,7 @@ def __enter__(self) -> LogfireSpan: self._span = self._tracer.start_span( name=self._span_name, attributes=self._otlp_attributes, + links=self._links, ) if self._token is None: # pragma: no branch self._token = context_api.attach(trace_api.set_span_in_context(self._span)) @@ -1864,6 +1872,12 @@ def set_attributes(self, attributes: dict[str, Any]) -> None: for key, value in attributes.items(): self.set_attribute(key, value) + def add_link(self, context: SpanContext, attributes: otel_types.Attributes = None) -> None: + if self._span is None: + self._links += [trace_api.Link(context=context, attributes=attributes)] + else: + self._span.add_link(context, attributes) + # TODO(Marcelo): We should add a test for `record_exception`. def record_exception( self, diff --git a/logfire/_internal/tracer.py b/logfire/_internal/tracer.py index 8ce2b8ab3..38b94be51 100644 --- a/logfire/_internal/tracer.py +++ b/logfire/_internal/tracer.py @@ -106,6 +106,9 @@ def set_attributes(self, attributes: dict[str, otel_types.AttributeValue]) -> No def set_attribute(self, key: str, value: otel_types.AttributeValue) -> None: self.span.set_attribute(key, value) + def add_link(self, context: SpanContext, attributes: otel_types.Attributes = None) -> None: + return self.span.add_link(context, attributes) + def add_event( self, name: str, diff --git a/tests/test_logfire.py b/tests/test_logfire.py index c581806e0..d2c27f360 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -2323,6 +2323,119 @@ def test_invalid_log_level(exporter: TestExporter): ) +def test_span_links(exporter: TestExporter): + with logfire.span('first span') as span: + first_context = span.context + + with logfire.span('second span') as span: + second_context = span.context + + assert first_context + assert second_context + with logfire.span('foo', _links=[(first_context, None)]) as span: + span.add_link(second_context) + + assert exporter.exported_spans_as_dict(_include_pending_spans=True)[-2:] == snapshot( + [ + { + 'name': 'foo (pending)', + 'context': {'trace_id': 3, 'span_id': 6, 'is_remote': False}, + 'parent': {'trace_id': 3, 'span_id': 5, 'is_remote': False}, + 'start_time': 5000000000, + 'end_time': 5000000000, + 'attributes': { + 'code.filepath': 'test_logfire.py', + 'code.function': 'test_span_links', + 'code.lineno': 123, + 'logfire.msg_template': 'foo', + 'logfire.msg': 'foo', + 'logfire.span_type': 'pending_span', + 'logfire.pending_parent_id': '0000000000000000', + }, + 'links': [{'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, 'attributes': {}}], + }, + { + 'name': 'foo', + 'context': {'trace_id': 3, 'span_id': 5, 'is_remote': False}, + 'parent': None, + 'start_time': 5000000000, + 'end_time': 6000000000, + 'attributes': { + 'code.filepath': 'test_logfire.py', + 'code.function': 'test_span_links', + 'code.lineno': 123, + 'logfire.msg_template': 'foo', + 'logfire.msg': 'foo', + 'logfire.span_type': 'span', + }, + 'links': [ + { + 'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, + 'attributes': {}, + }, + { + 'context': {'trace_id': 2, 'span_id': 3, 'is_remote': False}, + 'attributes': {}, + }, + ], + }, + ] + ) + + +def test_span_add_link_before_start(exporter: TestExporter): + with logfire.span('first span') as span: + context = span.context + + assert context + span = logfire.span('foo') + span.add_link(context) + + with span: + pass + + assert exporter.exported_spans_as_dict() == snapshot( + [ + { + 'name': 'first span', + 'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, + 'parent': None, + 'start_time': 1000000000, + 'end_time': 2000000000, + 'attributes': { + 'code.filepath': 'test_logfire.py', + 'code.function': 'test_span_add_link_before_start', + 'code.lineno': 123, + 'logfire.msg_template': 'first span', + 'logfire.msg': 'first span', + 'logfire.span_type': 'span', + }, + }, + { + 'name': 'foo', + 'context': {'trace_id': 2, 'span_id': 3, 'is_remote': False}, + 'parent': None, + 'start_time': 3000000000, + 'end_time': 4000000000, + 'attributes': { + 'code.filepath': 'test_logfire.py', + 'code.function': 'test_span_add_link_before_start', + 'code.lineno': 123, + 'logfire.msg_template': 'foo', + 'logfire.msg': 'foo', + 'logfire.span_type': 'span', + }, + 'links': [ + { + 'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, + 'attributes': {}, + } + ], + }, + ] + ) + + GLOBAL_VAR = 1