feat(tracer): add OpenTelemetryProvider implementation (#7003)

Author: Alwaysgaurav1

aws_lambda_powertools/tracing/__init__.py

@@ -1,6 +1,8 @@
 """Tracing utility"""
 
 from .extensions import aiohttp_trace_config
+from .opentelemetry import OpenTelemetryProvider, OpenTelemetrySegment
 from .tracer import Tracer
 
-__all__ = ["Tracer", "aiohttp_trace_config"]
+__all__ = ["OpenTelemetryProvider", "OpenTelemetrySegment", "Tracer", "aiohttp_trace_config"]
+

aws_lambda_powertools/tracing/opentelemetry.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Generator, Sequence
+
+from aws_lambda_powertools.tracing.base import BaseProvider, BaseSegment
+
+if TYPE_CHECKING:
+    import numbers
+    import traceback
+
+
+class OpenTelemetrySegment(BaseSegment):
+    """Segment implementation wrapping an OpenTelemetry Span."""
+
+    def __init__(self, span: Any):
+        self.span = span
+
+    def close(self, end_time: int | None = None):
+        if self.span and hasattr(self.span, "end"):
+            if end_time is not None:
+                self.span.end(end_time=int(end_time * 1e9))
+            else:
+                self.span.end()
+
+    def add_subsegment(self, subsegment: Any):
+        pass
+
+    def remove_subsegment(self, subsegment: Any):
+        pass
+
+    def put_annotation(self, key: str, value: str | numbers.Number | bool) -> None:
+        if self.span and hasattr(self.span, "set_attribute"):
+            self.span.set_attribute(key, value)
+
+    def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None:
+        if self.span and hasattr(self.span, "set_attribute"):
+            attr_key = f"{namespace}.{key}" if namespace else key
+            self.span.set_attribute(attr_key, str(value))
+
+    def add_exception(
+        self,
+        exception: BaseException,
+        stack: list[traceback.StackSummary] | None = None,
+        remote: bool = False,
+    ):
+        if self.span and hasattr(self.span, "record_exception"):
+            self.span.record_exception(exception)
+
+
+class OpenTelemetryProvider(BaseProvider):
+    """Tracing provider utilizing OpenTelemetry for Powertools Tracer."""
+
+    def __init__(self, tracer: Any | None = None):
+        if tracer is None:
+            try:
+                from opentelemetry import trace
+
+                tracer = trace.get_tracer("aws_lambda_powertools")
+            except ImportError:
+                tracer = None
+        self._tracer = tracer
+
+    @contextmanager
+    def in_subsegment(self, name: str | None = None, **kwargs) -> Generator[BaseSegment, None, None]:
+        name = name or "subsegment"
+        if self._tracer is not None:
+            with self._tracer.start_as_current_span(name) as span:
+                yield OpenTelemetrySegment(span)
+        else:
+            yield OpenTelemetrySegment(None)
+
+    @contextmanager
+    def in_subsegment_async(self, name: str | None = None, **kwargs) -> Generator[BaseSegment, None, None]:
+        name = name or "subsegment"
+        if self._tracer is not None:
+            with self._tracer.start_as_current_span(name) as span:
+                yield OpenTelemetrySegment(span)
+        else:
+            yield OpenTelemetrySegment(None)
+
+    def put_annotation(self, key: str, value: str | numbers.Number | bool) -> None:
+        try:
+            from opentelemetry import trace
+
+            span = trace.get_current_span()
+            if span and hasattr(span, "set_attribute"):
+                span.set_attribute(key, value)
+        except ImportError:
+            pass
+
+    def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None:
+        try:
+            from opentelemetry import trace
+
+            span = trace.get_current_span()
+            if span and hasattr(span, "set_attribute"):
+                attr_key = f"{namespace}.{key}" if namespace else key
+                span.set_attribute(attr_key, str(value))
+        except ImportError:
+            pass
+
+    def patch(self, modules: Sequence[str]) -> None:
+        pass
+
+    def patch_all(self) -> None:
+        pass

tests/unit/test_opentelemetry.py

@@ -0,0 +1,55 @@
+from unittest.mock import MagicMock
+
+from aws_lambda_powertools.tracing import OpenTelemetryProvider, OpenTelemetrySegment, Tracer
+
+
+def test_opentelemetry_segment_attributes():
+    mock_span = MagicMock()
+    segment = OpenTelemetrySegment(mock_span)
+
+    segment.put_annotation("key_ann", "val_ann")
+    mock_span.set_attribute.assert_called_with("key_ann", "val_ann")
+
+    segment.put_metadata("key_meta", {"data": 123}, namespace="test_ns")
+    mock_span.set_attribute.assert_called_with("test_ns.key_meta", "{'data': 123}")
+
+
+def test_opentelemetry_segment_exception():
+    mock_span = MagicMock()
+    segment = OpenTelemetrySegment(mock_span)
+    err = ValueError("test error")
+
+    segment.add_exception(err)
+    mock_span.record_exception.assert_called_with(err)
+
+
+def test_opentelemetry_provider_subsegment():
+    mock_tracer = MagicMock()
+    mock_span = MagicMock()
+    mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
+
+    provider = OpenTelemetryProvider(tracer=mock_tracer)
+
+    with provider.in_subsegment("my_subsegment") as sub:
+        assert isinstance(sub, OpenTelemetrySegment)
+        assert sub.span == mock_span
+
+    mock_tracer.start_as_current_span.assert_called_with("my_subsegment")
+
+
+def test_tracer_with_opentelemetry_provider():
+    mock_tracer = MagicMock()
+    mock_span = MagicMock()
+    mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
+
+    provider = OpenTelemetryProvider(tracer=mock_tracer)
+    tracer = Tracer(service="test-service", provider=provider, disabled=False)
+
+    assert tracer.provider == provider
+
+    @tracer.capture_method
+    def sample_func():
+        return "ok"
+
+    res = sample_func()
+    assert res == "ok"