first working POC for dd

Author: roger-zhangg

aws_lambda_powertools/tracing/dd_tracer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from numbers import Number
 from traceback import StackSummary
 from typing import Any, Generator, List, Optional, Sequence
@@ -8,12 +9,15 @@
 
 from .base import BaseProvider, BaseSegment
 
+logger = logging.getLogger(__name__)
+
 
 class DDSpan(BaseSegment):
     def __init__(self, dd_span=ddtrace.Span):
         self.dd_span = dd_span
 
     def close(self, end_time: int | None = None):
+        print("close is called")
         self.dd_span.finish(finish_time=float(end_time))
 
     def add_subsegment(self, subsegment: Any):
@@ -31,9 +35,22 @@ def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None
     def add_exception(self, exception: BaseException, stack: List[StackSummary], remote: bool = False):
         self.dd_span.set_exc_info(exc_type=exception, exc_val=exception, exc_tb=stack)
 
+    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)
+            print("exited")
+            self.close()
+        except Exception:
+            logger.exception("error closing trace")
+
 
 class DDTraceProvider(BaseProvider):
-    def __init__(self, dd_tracer=ddtrace.Tracer):
+    def __init__(self, dd_tracer: ddtrace.Tracer):
         self.dd_tracer = dd_tracer
 
     def in_subsegment(
@@ -42,18 +59,15 @@ def in_subsegment(
         service: Optional[str] = None,
         resource: Optional[str] = None,
         span_type: Optional[str] = None,
-        span_api: str = ddtrace.SPAN_API_DATADOG,
         **kwargs,
     ) -> Generator[BaseSegment, None, None]:
-        return self.dd_tracer.start_span(
+        dd_span = self.dd_tracer.trace(
             name,
-            child_of=self.dd_tracer.context_provider.active(),
             service=service,
             resource=resource,
             span_type=span_type,
-            activate=True,
-            span_api=span_api,
         )
+        return DDSpan(dd_span=dd_span)
 
     in_subsegment_async = in_subsegment