Merge pull request #56 from DataDog/tian.chu/fix-httplib-patch

Patching http clients erase existing headers

Author: tianchu

datadog_lambda/patch.py

@@ -10,27 +10,49 @@
 
 from wrapt import wrap_function_wrapper as wrap
 from wrapt.importer import when_imported
+from ddtrace import patch_all as patch_all_dd
 
-from datadog_lambda.tracing import get_dd_trace_context
+from datadog_lambda.tracing import (
+    get_dd_trace_context,
+    dd_tracing_enabled,
+)
 
 logger = logging.getLogger(__name__)
 
 if sys.version_info >= (3, 0, 0):
     httplib_module = "http.client"
+    from collections.abc import MutableMapping
 else:
     httplib_module = "httplib"
+    from collections import MutableMapping
 
 _httplib_patched = False
 _requests_patched = False
+_integration_tests_patched = False
 
 
 def patch_all():
     """
-    Patch the widely-used HTTP clients to automatically inject
-    Datadog trace context.
+    Patch third-party libraries for tracing.
     """
-    _patch_httplib()
-    _ensure_patch_requests()
+    _patch_for_integration_tests()
+
+    if dd_tracing_enabled:
+        patch_all_dd()
+    else:
+        _patch_httplib()
+        _ensure_patch_requests()
+
+
+def _patch_for_integration_tests():
+    """
+    Patch `requests` to log the outgoing requests for integration tests.
+    """
+    global _integration_tests_patched
+    is_in_tests = os.environ.get("DD_INTEGRATION_TEST", "false").lower() == "true"
+    if not _integration_tests_patched and is_in_tests:
+        wrap("requests", "Session.send", _log_request)
+        _integration_tests_patched = True
 
 
 def _patch_httplib():
@@ -80,17 +102,13 @@ def _wrap_requests_request(func, instance, args, kwargs):
     into the outgoing requests.
     """
     context = get_dd_trace_context()
-    if "headers" in kwargs and isinstance(kwargs["headers"], dict):
+    if "headers" in kwargs and isinstance(kwargs["headers"], MutableMapping):
         kwargs["headers"].update(context)
-    elif len(args) >= 5 and isinstance(args[4], dict):
+    elif len(args) >= 5 and isinstance(args[4], MutableMapping):
         args[4].update(context)
     else:
         kwargs["headers"] = context
 
-    # If we're in an integration test, log the HTTP requests made
-    if os.environ.get("DD_INTEGRATION_TEST", "false").lower() == "true":
-        _print_request_string(args, kwargs)
-
     return func(*args, **kwargs)
 
 
@@ -100,43 +118,38 @@ def _wrap_httplib_request(func, instance, args, kwargs):
     the Datadog trace headers into the outgoing requests.
     """
     context = get_dd_trace_context()
-    if "headers" in kwargs and isinstance(kwargs["headers"], dict):
+    if "headers" in kwargs and isinstance(kwargs["headers"], MutableMapping):
         kwargs["headers"].update(context)
-    elif len(args) >= 4 and isinstance(args[3], dict):
+    elif len(args) >= 4 and isinstance(args[3], MutableMapping):
         args[3].update(context)
     else:
         kwargs["headers"] = context
 
     return func(*args, **kwargs)
 
 
-def _print_request_string(args, kwargs):
+def _log_request(func, instance, args, kwargs):
+    request = kwargs.get("request") or args[0]
+    _print_request_string(request)
+    return func(*args, **kwargs)
+
+
+def _print_request_string(request):
     """Print the request so that it can be checked in integration tests
 
     Only used by integration tests.
     """
-    # Normalizes the different ways args can be passed to a request
-    # to prevent test flakiness
-    method = None
-    if len(args) > 0:
-        method = args[0]
-    else:
-        method = kwargs.get("method", "").upper()
-
-    url = None
-    if len(args) > 1:
-        url = args[1]
-    else:
-        url = kwargs.get("url")
+    method = request.method
+    url = request.url
 
     # Sort the datapoints POSTed by their name so that snapshots always align
-    data = kwargs.get("data", "{}")
+    data = request.body or "{}"
     data_dict = json.loads(data)
     data_dict.get("series", []).sort(key=lambda series: series.get("metric"))
     sorted_data = json.dumps(data_dict)
 
     # Sort headers to prevent any differences in ordering
-    headers = kwargs.get("headers", {})
+    headers = request.headers or {}
     sorted_headers = sorted(
         "{}:{}".format(key, value) for key, value in headers.items()
     )

datadog_lambda/wrapper.py

@@ -22,7 +22,6 @@
     set_dd_trace_py_root,
     create_function_execution_span,
 )
-from ddtrace import patch_all as patch_all_dd
 
 logger = logging.getLogger(__name__)
 
@@ -93,12 +92,9 @@ def __init__(self, func):
             if self.logs_injection:
                 inject_correlation_ids()
 
-            if not dd_tracing_enabled:
-                # When using dd_trace_py it will patch all the http clients for us,
-                # Patch HTTP clients to propagate Datadog trace context
-                patch_all()
-            else:
-                patch_all_dd()
+            # Patch third-party libraries for tracing
+            patch_all()
+
             logger.debug("datadog_lambda_wrapper initialized")
         except Exception:
             traceback.print_exc()
@@ -137,10 +133,10 @@ def _before(self, event, context):
 
     def _after(self, event, context):
         try:
-            if self.span:
-                self.span.finish()
             if not self.flush_to_log:
                 lambda_stats.flush(float("inf"))
+            if self.span:
+                self.span.finish()
             logger.debug("datadog_lambda_wrapper _after() done")
         except Exception:
             traceback.print_exc()