Add start, end, attach, and detach to LogfireSpan (#914)

Author: alexmojaki

logfire/_internal/main.py

@@ -2138,35 +2138,48 @@ def __init__(
         def __getattr__(self, name: str) -> Any:
             return getattr(self._span, name)
 
-    def __enter__(self) -> LogfireSpan:
-        with handle_internal_errors:
-            if self._span is None:  # pragma: no branch
-                self._span = self._tracer.start_span(
-                    name=self._span_name,
-                    attributes=self._otlp_attributes,
-                    links=self._links,
-                )
-            self._span.__enter__()
-            if self._token is None:  # pragma: no branch
-                self._token = context_api.attach(trace_api.set_span_in_context(self._span))
+    @handle_internal_errors
+    def _start(self):
+        if self._span is not None:
+            return
+        self._span = self._tracer.start_span(
+            name=self._span_name,
+            attributes=self._otlp_attributes,
+            links=self._links,
+        )
+
+    @handle_internal_errors
+    def _attach(self):
+        if self._token is not None:
+            return
+        assert self._span is not None
+        self._token = context_api.attach(trace_api.set_span_in_context(self._span))
 
+    def __enter__(self) -> LogfireSpan:
+        self._start()
+        self._attach()
         return self
 
     @handle_internal_errors
-    def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None:
-        if self._token is None:  # pragma: no cover
+    def _end(self):
+        if not self._span or not self._span.is_recording():
             return
-        assert self._span is not None
+        if self._added_attributes:
+            self._span.set_attribute(ATTRIBUTES_JSON_SCHEMA_KEY, attributes_json_schema(self._json_schema_properties))
+        self._span.end()
 
+    def _detach(self):
+        if self._token is None:
+            return
         context_api.detach(self._token)
         self._token = None
-        if self._span.is_recording():
-            with handle_internal_errors:
-                if self._added_attributes:
-                    self._span.set_attribute(
-                        ATTRIBUTES_JSON_SCHEMA_KEY, attributes_json_schema(self._json_schema_properties)
-                    )
-            self._span.__exit__(exc_type, exc_value, traceback)
+
+    @handle_internal_errors
+    def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None:
+        self._detach()
+        if self._span and self._span.is_recording() and isinstance(exc_value, BaseException):
+            self._span.record_exception(exc_value, escaped=True)
+        self._end()
 
     @property
     def message_template(self) -> str | None:  # pragma: no cover

tests/test_logfire.py

@@ -19,7 +19,7 @@
 from opentelemetry.sdk.metrics.export import InMemoryMetricReader
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
-from opentelemetry.trace import StatusCode, get_tracer
+from opentelemetry.trace import INVALID_SPAN, StatusCode, get_current_span, get_tracer
 from pydantic import BaseModel, __version__ as pydantic_version
 from pydantic_core import ValidationError
 
@@ -3323,3 +3323,70 @@ def test_default_id_generator(exporter: TestExporter) -> None:
         export['attributes']['i'] for export in sorted(exported, key=lambda span: span['start_time'])
     ]
     assert sorted_by_trace_id == sorted_by_start_timestamp
+
+
+def test_start_end_attach_detach(exporter: TestExporter, caplog: pytest.LogCaptureFixture):
+    span = logfire.span('test')
+    assert not span.is_recording()
+    assert not exporter.exported_spans
+    assert get_current_span() is INVALID_SPAN
+
+    span._start()  # type: ignore
+    span._start()  # type: ignore
+    assert span.is_recording()
+    assert len(exporter.exported_spans) == 1
+    assert get_current_span() is INVALID_SPAN
+
+    span._attach()  # type: ignore
+    span._attach()  # type: ignore
+    assert get_current_span() is span._span  # type: ignore
+
+    span._detach()  # type: ignore
+    span._detach()  # type: ignore
+    assert span.is_recording()
+    assert len(exporter.exported_spans) == 1
+    assert get_current_span() is INVALID_SPAN
+
+    span._end()  # type: ignore
+    span._end()  # type: ignore
+    assert not span.is_recording()
+    assert len(exporter.exported_spans) == 2
+    assert get_current_span() is INVALID_SPAN
+
+    assert not caplog.messages
+
+    assert exporter.exported_spans_as_dict(_include_pending_spans=True) == snapshot(
+        [
+            {
+                'name': 'test',
+                'context': {'trace_id': 1, 'span_id': 2, 'is_remote': False},
+                'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'start_time': 1000000000,
+                'end_time': 1000000000,
+                'attributes': {
+                    'code.filepath': 'test_logfire.py',
+                    'code.function': 'test_start_end_attach_detach',
+                    'code.lineno': 123,
+                    'logfire.msg_template': 'test',
+                    'logfire.msg': 'test',
+                    'logfire.span_type': 'pending_span',
+                    'logfire.pending_parent_id': '0000000000000000',
+                },
+            },
+            {
+                'name': 'test',
+                '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_start_end_attach_detach',
+                    'code.lineno': 123,
+                    'logfire.msg_template': 'test',
+                    'logfire.msg': 'test',
+                    'logfire.span_type': 'span',
+                },
+            },
+        ]
+    )