fix mypy, move into folder

Author: roger-zhangg

aws_lambda_powertools/tracing/provider/__init__.py

@@ -1,4 +1,5 @@
-from .dd_tracer import DDTraceProvider
-from .otel_tracer import OtelProvider
+from .datadog.dd_tracer import DDTraceProvider
+from .otel.otel_tracer import OtelProvider
+from .xray.xray_tracer import XrayProvider
 
-__all__ = ["DDTraceProvider", "OtelProvider"]
+__all__ = ["DDTraceProvider", "OtelProvider", "XrayProvider"]

aws_lambda_powertools/tracing/provider/base.py

@@ -1,15 +1,13 @@
 import abc
 import numbers
-from contextlib import contextmanager
-from typing import Generator, Sequence, Union
+from contextlib import asynccontextmanager, contextmanager
+from typing import AsyncGenerator, Generator, Sequence, Union
 
 
 ## TO-Discuss how to refactor this one. Segment seems X-ray exclusive concept
 class BaseSpan(abc.ABC):
     """Holds common properties and methods on segment and subsegment."""
 
-    """Triggering"""
-
     @abc.abstractmethod
     def set_attribute(self, key: str, value: Union[str, numbers.Number, bool], **kwargs) -> None:
         """Annotate segment or subsegment with a key-value pair.
@@ -39,7 +37,7 @@ def record_exception(self, exception: BaseException, **kwargs):
 class BaseProvider(abc.ABC):
     @abc.abstractmethod
     @contextmanager
-    def trace(self, name=None, **kwargs) -> Generator[BaseSpan, None, None]:
+    def trace(self, name: str, **kwargs) -> Generator[BaseSpan, None, None]:
         """Return a span context manger.
 
         Parameters
@@ -51,8 +49,8 @@ def trace(self, name=None, **kwargs) -> Generator[BaseSpan, None, None]:
         """
 
     @abc.abstractmethod
-    @contextmanager
-    def trace_async(self, name=None, **kwargs) -> Generator[BaseSpan, None, None]:
+    @asynccontextmanager
+    def trace_async(self, name: str, **kwargs) -> AsyncGenerator[BaseSpan, None]:
         """Return a async span context manger.
 
         Parameters

aws_lambda_powertools/tracing/provider/datadog/dd_tracer.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import logging
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, AsyncGenerator, Dict, Generator, Optional, Sequence
+
+import ddtrace
+
+from ..base import BaseProvider, BaseSpan
+
+logger = logging.getLogger(__name__)
+
+
+class DDSpan(BaseSpan):
+    def __init__(self, dd_span=ddtrace.Span):
+        self.dd_span = dd_span
+
+    def set_attribute(self, key: str | bytes, value: Any, **kwargs) -> None:
+        self.dd_span.set_tag(key=key, value=value)
+
+    def record_exception(self, exception: BaseException, attributes: Optional[Dict] = None, **kwargs):
+        self.dd_span.set_traceback()
+        _attributes = attributes or {}
+        if kwargs:
+            kwargs.update(_attributes)
+            # attribute should overwrite kwargs
+            _attributes = kwargs
+        self.dd_span.set_tags(tags=attributes)
+
+    def __enter__(self):
+        print("entered")
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        try:
+            if exc_type:
+                self.dd_span.set_exc_info(exc_type, exc_val, exc_tb)
+            self.dd_span.finish()
+        except Exception as e:
+            logger.exception(f"error closing trace {e}")
+
+
+class DDTraceProvider(BaseProvider):
+    def __init__(self, dd_tracer: ddtrace.Tracer):
+        self.dd_tracer = dd_tracer
+
+    @contextmanager
+    def trace(
+        self,
+        name: str,
+        service: Optional[str] = None,
+        resource: Optional[str] = None,
+        span_type: Optional[str] = None,
+        **kwargs,
+    ) -> Generator[DDSpan, None, None]:
+        dd_span = self.dd_tracer.trace(
+            name=name,
+            service=service,
+            resource=resource,
+            span_type=span_type,
+        )
+        yield DDSpan(dd_span=dd_span)
+
+    in_subsegment = trace
+
+    @asynccontextmanager
+    async def trace_async(
+        self,
+        name: str,
+        service: Optional[str] = None,
+        resource: Optional[str] = None,
+        span_type: Optional[str] = None,
+        **kwargs,
+    ) -> AsyncGenerator[DDSpan, None]:
+        dd_span = self.dd_tracer.trace(
+            name=name,
+            service=service,
+            resource=resource,
+            span_type=span_type,
+        )
+        yield DDSpan(dd_span=dd_span)
+
+    def set_attribute(self, key: str | bytes, value: Any, **kwargs: Any) -> None:
+        span = self.dd_tracer.context_provider.active()
+        if isinstance(span, ddtrace.Span):
+            span.set_tag(key=key, value=value)
+        # ignore if no active span
+
+    def patch(self, modules: Sequence[str]) -> None:
+        module_to_patch = {}
+        for m in modules:
+            module_to_patch[m] = True
+        ddtrace.patch(**module_to_patch)  # type:ignore
+
+    def patch_all(self) -> None:
+        ddtrace.patch_all()

aws_lambda_powertools/tracing/provider/dd_tracer.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import logging
+from contextlib import asynccontextmanager, contextmanager
+from typing import AsyncGenerator, Dict, Generator, Optional, Sequence
+
+from opentelemetry import trace as otel_trace
+
+from ..base import BaseProvider, BaseSpan
+
+logger = logging.getLogger(__name__)
+
+
+# optl terminology first
+# 1. Provider based on OTel terminology
+# 2. X-Ray provider on top of the new BaseProvider
+# 3. Datadog provider on top of the new BaseProvider
+# access xray sdk
+class OtelSpan(BaseSpan):
+    def __init__(self, otel_span=otel_trace.Span):
+        self.otel_span = otel_span
+
+    def set_attribute(self, key: str, value: str | float | bool, **kwargs) -> None:  # type: ignore[override]
+        self.otel_span.set_attribute(key=key, value=value)
+
+    def record_exception(
+        self,
+        exception: BaseException,
+        attributes: Optional[Dict] = None,
+        timestamp: Optional[int] = None,
+        escaped: bool = False,
+        **kwargs,
+    ):
+        _attributes = attributes or {}
+        if kwargs:
+            kwargs.update(_attributes)
+            # attribute should overwrite kwargs
+            _attributes = kwargs
+        self.otel_span.record_exception(
+            exception=exception,
+            attributes=_attributes,
+            timestamp=timestamp,
+            escaped=escaped,
+        )
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        try:
+            self.otel_span.__exit__(exc_type, exc_val, exc_tb)
+        except Exception as e:
+            logger.exception(f"error closing trace {e}")
+
+
+class OtelProvider(BaseProvider):
+    def __init__(self, otel_tracer: otel_trace.Tracer):
+        self.otel_tracer = otel_tracer
+
+    @contextmanager
+    def trace(self, name: str, **kwargs) -> Generator[OtelSpan, None, None]:
+        with self.otel_tracer.start_as_current_span(name=name, **kwargs) as otel_span:
+            yield OtelSpan(otel_span=otel_span)
+
+    in_subsegment = trace
+
+    @asynccontextmanager
+    async def trace_async(self, name: str, **kwargs) -> AsyncGenerator[OtelSpan, None]:
+        with self.otel_tracer.start_as_current_span(name=name, **kwargs) as otel_span:
+            yield OtelSpan(otel_span=otel_span)
+
+    in_subsegment_async = trace_async
+
+    def set_attribute(self, key: str, value: str | float | bool, **kwargs) -> None:  # type: ignore[override]
+        active_span = otel_trace.get_current_span()
+        active_span.set_attribute(key=key, value=value)
+
+    def patch(self, modules: Sequence[str]) -> None:
+        # OTEL sdk doesn't have patch
+        pass
+
+    def patch_all(self) -> None:
+        pass