Skip to content

Commit

Permalink
Add span links (#587)
Browse files Browse the repository at this point in the history 
Co-authored-by: Alex Hall <[email protected]>
  • Loading branch information
@Kludex @alexmojaki
Kludex and alexmojaki authored Nov 14, 2024
1 parent 183b3b2 commit 9001475
Show file tree
Hide file tree
Showing 4 changed files with 143 additions and 17 deletions.
28 changes: 12 additions & 16 deletions logfire/_internal/exporters/test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
16 changes: 15 additions & 1 deletion logfire/_internal/main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -200,6 +201,7 @@ def _span(
otlp_attributes,
self._spans_tracer,
json_schema_properties,
links=_links,
)
except Exception:
log_internal_error()
Expand Down Expand Up @@ -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.
Expand All @@ -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.
"""
Expand All @@ -521,6 +525,7 @@ def span(
_tags=_tags,
_span_name=_span_name,
_level=_level,
_links=_links,
)

def instrument(
Expand Down Expand Up @@ -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
Expand All @@ -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))
Expand Down Expand Up @@ -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,
Expand Down
3 changes: 3 additions & 0 deletions logfire/_internal/tracer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
113 changes: 113 additions & 0 deletions tests/test_logfire.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down

0 comments on commit 9001475

Please sign in to comment.