Fix issues with reading trace context from ddtrace

Author: DarcyRaynerDD

datadog_lambda/trace_wrapper.py

@@ -22,7 +22,7 @@ def start_span(self, name, **kwargs):
         if not tracer:
             return None
 
-        return self._tracer.start_span(name, **kwargs)
+        return tracer.start_span(name, **kwargs)
 
     def _load_tracer(self):
         if not TraceWrapper.tracer_enabled():
@@ -50,8 +50,8 @@ def trace_context(self):
         parent_id = span.context.span_id
         trace_id = span.context.trace_id
         return {
-            "parent_id": parent_id,
-            "trace_id": trace_id,
+            "parent_id": str(parent_id),
+            "trace_id": str(trace_id),
             "sampling_priority": SamplingPriority.AUTO_KEEP,
             "source": Source.DDTRACE,
         }

datadog_lambda/tracing.py

@@ -62,9 +62,9 @@ def _get_xray_trace_context():
 
 def _context_obj_to_headers(obj):
     return {
-        TraceHeader.TRACE_ID: obj.get("trace_id"),
-        TraceHeader.PARENT_ID: obj.get("parent_id"),
-        TraceHeader.SAMPLING_PRIORITY: obj.get("sampling_priority"),
+        TraceHeader.TRACE_ID: str(obj.get("trace_id")),
+        TraceHeader.PARENT_ID: str(obj.get("parent_id")),
+        TraceHeader.SAMPLING_PRIORITY: str(obj.get("sampling_priority")),
     }
 
 
@@ -122,23 +122,26 @@ def get_dd_trace_context():
     """
     global dd_trace_context
 
-    if not dd_trace_context:
-        return None
-    trace_context = _context_obj_to_headers(dd_trace_context)
-    datadog_context = trace_wrapper.trace_context
-    if datadog_context:
-        logger.debug("get_dd_trace_context using dd-trace context")
-        return datadog_context
+    native_trace_context = trace_wrapper.trace_context
+    if native_trace_context:
+        logger.info("get_dd_trace_context using dd-trace context")
+        return _context_obj_to_headers(native_trace_context)
+
     try:
+        trace_headers = _context_obj_to_headers(dd_trace_context)
         xray_context = _get_xray_trace_context()  # xray (sub)segment
-        if xray_context:
-            trace_context[TraceHeader.PARENT_ID] = xray_context["parent_id"]
+        if xray_context and not trace_headers:
+            return _context_obj_to_headers(xray_context)
+        if xray_context and trace_headers:
+            trace_headers[TraceHeader.PARENT_ID] = xray_context["parent-id"]
+            return trace_headers
     except Exception as e:
         logger.debug(
             "get_dd_trace_context couldn't read from segment from x-ray, with error %s"
             % e
         )
-    return trace_context
+
+    return {}
 
 
 def set_correlation_ids():

setup.py

@@ -6,40 +6,34 @@
 
 here = path.abspath(path.dirname(__file__))
 
-with open(path.join(here, 'README.md'), encoding='utf-8') as f:
+with open(path.join(here, "README.md"), encoding="utf-8") as f:
     long_description = f.read()
 
 setup(
-    name='datadog_lambda',
+    name="datadog_lambda",
     version=__version__,
-    description='The Datadog AWS Lambda Layer',
+    description="The Datadog AWS Lambda Layer",
     long_description=long_description,
-    long_description_content_type='text/markdown',
-    url='https://github.com/DataDog/datadog-lambda-layer-python',
-    author='Datadog, Inc.',
-    author_email='dev@datadoghq.com',
+    long_description_content_type="text/markdown",
+    url="https://github.com/DataDog/datadog-lambda-layer-python",
+    author="Datadog, Inc.",
+    author_email="dev@datadoghq.com",
     classifiers=[
-        'Programming Language :: Python :: 2.7',
-        'Programming Language :: Python :: 3.6',
-        'Programming Language :: Python :: 3.7',
-        'Programming Language :: Python :: 3.8',
+        "Programming Language :: Python :: 2.7",
+        "Programming Language :: Python :: 3.6",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
     ],
-    keywords='datadog aws lambda layer',
-    packages=['datadog_lambda'],
-    python_requires='>=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, <4',
+    keywords="datadog aws lambda layer",
+    packages=["datadog_lambda"],
+    python_requires=">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, <4",
     install_requires=[
-        'aws-xray-sdk==2.4.3',
-        'datadog==0.32.0',
-        'ddtrace==0.31.0',
-        'wrapt==1.11.2',
-        'setuptools==42.0.2',
+        "aws-xray-sdk==2.4.3",
+        "datadog==0.32.0",
+        "wrapt==1.11.2",
+        "setuptools==42.0.2",
     ],
     extras_require={
-        'dev': [
-            'nose2==0.9.1',
-            'flake8==3.7.9',
-            'requests==2.22.0',
-            'boto3==1.10.33',
-        ]
-    }
+        "dev": ["nose2==0.9.1", "flake8==3.7.9", "requests==2.22.0", "boto3==1.10.33"]
+    },
 )