init laminar span context from env var (#221)

* init laminar span context from env var

* small fixes and unit tests

* please the bots in tests

* bump version to 0.7.24

---------

Co-authored-by: Din <[email protected]>

Author: Rainhunter13

pyproject.toml

@@ -6,7 +6,7 @@
 
 [project]
 name = "lmnr"
-version = "0.7.23"
+version = "0.7.24"
 description = "Python SDK for Laminar"
 authors = [
   { name = "lmnr.ai", email = "[email protected]" }

src/lmnr/opentelemetry_lib/tracing/processor.py

@@ -140,3 +140,13 @@ def clear(self):
         with self._paths_lock:
             self.__span_id_to_path = {}
             self.__span_id_lists = {}
+
+    def set_parent_path_info(
+        self,
+        parent_span_id: int,
+        span_path: list[str],
+        span_ids_path: list[str],
+    ):
+        with self._paths_lock:
+            self.__span_id_to_path[parent_span_id] = span_path
+            self.__span_id_lists[parent_span_id] = span_ids_path

src/lmnr/sdk/laminar.py

@@ -8,8 +8,10 @@
     CONTEXT_USER_ID_KEY,
     attach_context,
     get_event_attributes_from_context,
+    push_span_context,
 )
 from lmnr.opentelemetry_lib.tracing.instruments import Instruments
+from lmnr.opentelemetry_lib.tracing.processor import LaminarSpanProcessor
 from lmnr.opentelemetry_lib.tracing.tracer import get_tracer_with_context
 from lmnr.opentelemetry_lib.tracing.attributes import (
     ASSOCIATION_PROPERTIES,
@@ -203,6 +205,46 @@ def initialize(
             force_http=force_http,
         )
 
+        cls._initialize_context_from_env()
+
+    @classmethod
+    def _initialize_context_from_env(cls) -> None:
+        """Attach upstream Laminar context from the environment, if provided."""
+        env_context = os.getenv("LMNR_SPAN_CONTEXT")
+        if not env_context:
+            return
+
+        try:
+            laminar_context = LaminarSpanContext.deserialize(env_context)
+        except Exception as exc:  # pylint: disable=broad-exception-caught
+            cls.__logger.warning(
+                "LMNR_SPAN_CONTEXT is set but could not be deserialized: %s", exc
+            )
+            return
+
+        try:
+            otel_span_context = LaminarSpanContext.try_to_otel_span_context(
+                laminar_context, cls.__logger
+            )
+        except ValueError as exc:
+            cls.__logger.warning(
+                "LMNR_SPAN_CONTEXT is set but invalid span context provided: %s", exc
+            )
+            return
+
+        base_context = trace.set_span_in_context(
+            trace.NonRecordingSpan(otel_span_context), get_current_context()
+        )
+        processor = TracerWrapper.instance._span_processor
+        if isinstance(processor, LaminarSpanProcessor):
+            processor.set_parent_path_info(
+                otel_span_context.span_id,
+                laminar_context.span_path,
+                laminar_context.span_ids_path,
+            )
+        push_span_context(base_context)
+        cls.__logger.debug("Initialized Laminar parent context from LMNR_SPAN_CONTEXT.")
+
     @classmethod
     def is_initialized(cls):
         """Check if Laminar is initialized. A utility to make sure other

src/lmnr/version.py

@@ -3,7 +3,7 @@
 from packaging import version
 
 
-__version__ = "0.7.23"
+__version__ = "0.7.24"
 PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"
 
 

tests/test_observe.py

@@ -1,8 +1,9 @@
 import json
+import os
 import uuid
 import pytest
 
-from lmnr import Laminar, observe
+from lmnr import Laminar, observe, LaminarSpanContext
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from opentelemetry import trace
 
@@ -1563,3 +1564,48 @@ def func2():
         str(uuid.UUID(int=func1_span.get_span_context().span_id)),
         str(uuid.UUID(int=func2_span.get_span_context().span_id)),
     )
+
+
+def test_span_context_from_env_variables_observe(span_exporter: InMemorySpanExporter):
+    test_trace_id = "01234567-89ab-cdef-0123-456789abcdef"
+    test_span_id = "00000000-0000-0000-0123-456789abcdef"
+    test_span_id2 = "00000000-0000-0000-fedc-ba9876543210"
+    old_val = os.getenv("LMNR_SPAN_CONTEXT")
+    test_context = LaminarSpanContext(
+        trace_id=test_trace_id,
+        span_id=test_span_id2,
+        span_path=["grandparent", "parent"],
+        span_ids_path=[test_span_id, test_span_id2],
+    )
+
+    os.environ["LMNR_SPAN_CONTEXT"] = str(test_context)
+
+    Laminar._initialize_context_from_env()
+
+    @observe()
+    def test():
+        pass
+
+    test()
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span_id = spans[0].get_span_context().span_id
+    assert spans[0].name == "test"
+    assert spans[0].attributes["lmnr.span.instrumentation_source"] == "python"
+    assert spans[0].attributes["lmnr.span.path"] == (
+        "grandparent",
+        "parent",
+        "test",
+    )
+    assert spans[0].attributes["lmnr.span.ids_path"] == (
+        str(uuid.UUID(test_span_id)),
+        str(uuid.UUID(test_span_id2)),
+        str(uuid.UUID(int=span_id)),
+    )
+    assert spans[0].get_span_context().trace_id == uuid.UUID(test_trace_id).int
+    assert spans[0].parent.span_id == uuid.UUID(test_span_id2).int
+    if old_val:
+        os.environ["LMNR_SPAN_CONTEXT"] = old_val
+    else:
+        os.environ.pop("LMNR_SPAN_CONTEXT", None)

tests/test_tracing.py

@@ -1,4 +1,5 @@
 import json
+import os
 import pytest
 import uuid
 
@@ -671,6 +672,47 @@ def foo(context: LaminarSpanContext):
     assert inner_span.parent.span_id == outer_span.get_span_context().span_id
 
 
+def test_span_context_from_env_variables(span_exporter: InMemorySpanExporter):
+    test_trace_id = "01234567-89ab-cdef-0123-456789abcdef"
+    test_span_id = "00000000-0000-0000-0123-456789abcdef"
+    test_span_id2 = "00000000-0000-0000-fedc-ba9876543210"
+    old_val = os.getenv("LMNR_SPAN_CONTEXT")
+    test_context = LaminarSpanContext(
+        trace_id=test_trace_id,
+        span_id=test_span_id2,
+        span_path=["grandparent", "parent"],
+        span_ids_path=[test_span_id, test_span_id2],
+    )
+
+    os.environ["LMNR_SPAN_CONTEXT"] = str(test_context)
+
+    Laminar._initialize_context_from_env()
+    with Laminar.start_as_current_span("test"):
+        pass
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span_id = spans[0].get_span_context().span_id
+    assert spans[0].name == "test"
+    assert spans[0].attributes["lmnr.span.instrumentation_source"] == "python"
+    assert spans[0].attributes["lmnr.span.path"] == (
+        "grandparent",
+        "parent",
+        "test",
+    )
+    assert spans[0].attributes["lmnr.span.ids_path"] == (
+        str(uuid.UUID(test_span_id)),
+        str(uuid.UUID(test_span_id2)),
+        str(uuid.UUID(int=span_id)),
+    )
+    assert spans[0].get_span_context().trace_id == uuid.UUID(test_trace_id).int
+    assert spans[0].parent.span_id == uuid.UUID(test_span_id2).int
+    if old_val:
+        os.environ["LMNR_SPAN_CONTEXT"] = old_val
+    else:
+        os.environ.pop("LMNR_SPAN_CONTEXT", None)
+
+
 def test_tags_deduplication(span_exporter: InMemorySpanExporter):
     with Laminar.start_as_current_span("test"):
         Laminar.set_span_tags(["foo", "bar", "foo"])