Implement PR feedback

Author: DarcyRaynerDD

datadog_lambda/tracing.py

@@ -16,12 +16,15 @@
     TraceContextSource,
 )
 from ddtrace import tracer, patch
+from ddtrace.propagation.http import HTTPPropagator
 
 logger = logging.getLogger(__name__)
 
 dd_trace_context = {}
 dd_tracing_enabled = os.environ.get("DD_TRACE_ENABLED", "false").lower() == "true"
 
+propagator = HTTPPropagator()
+
 
 def _convert_xray_trace_id(xray_trace_id):
     """
@@ -221,3 +224,36 @@ def is_lambda_context():
     regular `Context` (e.g., when testing lambda functions locally).
     """
     return type(xray_recorder.context) == LambdaContext
+
+
+def create_function_execution_span(
+    context,
+    function_name,
+    handler_name,
+    is_cold_start,
+    trace_context,
+    merge_xray_traces,
+):
+    span_context = None
+    if trace_context["source"] == TraceContextSource.EVENT or merge_xray_traces:
+        headers = get_dd_trace_context()
+        span_context = propagator.extract(headers)
+
+    tags = {}
+    if context:
+        tags = {
+            "cold_start": is_cold_start,
+            "function_arn": context.invoked_function_arn,
+            "request_id": context.aws_request_id,
+            "resource_names": context.function_name,
+        }
+    args = {
+        "service": function_name,
+        "resource": handler_name,
+        "span_type": "serverless",
+        "child_of": span_context,
+    }
+    span = tracer.start_span("aws.lambda", **args)
+    if span:
+        span.set_tags(tags)
+    return span

datadog_lambda/wrapper.py

@@ -17,14 +17,10 @@
 from datadog_lambda.tracing import (
     extract_dd_trace_context,
     inject_correlation_ids,
-    get_dd_trace_context,
     dd_tracing_enabled,
     set_correlation_ids,
+    create_function_execution_span,
 )
-from datadog_lambda.constants import TraceContextSource
-from ddtrace import tracer
-from ddtrace.propagation.http import HTTPPropagator
-
 
 logger = logging.getLogger(__name__)
 
@@ -91,7 +87,6 @@ def __init__(self, func):
             )
             self.handler_name = os.environ.get("_HANDLER", "handler")
             self.function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME", "function")
-            self.propagator = HTTPPropagator()
 
             # Inject trace correlation ids to logs
             if self.logs_injection:
@@ -126,7 +121,15 @@ def _before(self, event, context):
 
             self.span = None
             if dd_tracing_enabled:
-                self.span = self._create_dd_trace_py_span(context, dd_context)
+
+                self.span = create_function_execution_span(
+                    context,
+                    self.function_name,
+                    self.handler_name,
+                    is_cold_start(),
+                    dd_context,
+                    self.merge_xray_traces,
+                )
             else:
                 set_correlation_ids()
 
@@ -144,33 +147,5 @@ def _after(self, event, context):
         except Exception:
             traceback.print_exc()
 
-    def _create_dd_trace_py_span(self, context, trace_context):
-        span_context = None
-        if (
-            trace_context["source"] == TraceContextSource.EVENT
-            or self.merge_xray_traces
-        ):
-            headers = get_dd_trace_context()
-            span_context = self.propagator.extract(headers)
-
-        tags = {}
-        if context:
-            tags = {
-                "cold_start": is_cold_start(),
-                "function_arn": context.invoked_function_arn,
-                "request_id": context.aws_request_id,
-                "resource_names": context.function_name,
-            }
-        args = {
-            "service": self.function_name,
-            "resource": self.handler_name,
-            "span_type": "serverless",
-            "child_of": span_context,
-        }
-        span = tracer.start_span("aws.lambda", **args)
-        if span:
-            span.set_tags(tags)
-        return span
-
 
 datadog_lambda_wrapper = _LambdaDecorator

tests/Dockerfile

@@ -8,5 +8,4 @@ WORKDIR /test
 
 # Install datadog-lambda with dev dependencies from local
 COPY . .
-RUN pip install .[dev]
-RUN pip install --find-links=https://s3.amazonaws.com/pypi.datadoghq.com/trace-dev/index.html ddtrace==0.35.1.dev5+g1e11b2dd -t .[dev]
+RUN pip install .[dev]