patch requests when it's used

Author: tianchu

datadog_lambda/patch.py

@@ -9,6 +9,7 @@
 import logging
 
 from wrapt import wrap_function_wrapper as wrap
+from wrapt.importer import when_imported
 
 from datadog_lambda.tracing import get_dd_trace_context
 
@@ -29,7 +30,7 @@ def patch_all():
     Datadog trace context.
     """
     _patch_httplib()
-    _patch_requests()
+    _ensure_patch_requests()
 
 
 def _patch_httplib():
@@ -45,7 +46,20 @@ def _patch_httplib():
     logger.debug("Patched %s", httplib_module)
 
 
-def _patch_requests():
+def _ensure_patch_requests():
+    """
+    `requests` is third-party, may not be installed or used,
+    but ensure it gets patched if installed and used.
+    """
+    if "requests" in sys.modules:
+        # already imported, patch now
+        _patch_requests(sys.modules["requests"])
+    else:
+        # patch when imported
+        when_imported("requests")(_patch_requests)
+
+
+def _patch_requests(module):
     """
     Patch the high-level HTTP client module `requests`
     if it's installed.
@@ -66,9 +80,9 @@ def _wrap_requests_request(func, instance, args, kwargs):
     into the outgoing requests.
     """
     context = get_dd_trace_context()
-    if "headers" in kwargs:
+    if "headers" in kwargs and isinstance(kwargs["headers"], dict):
         kwargs["headers"].update(context)
-    elif len(args) >= 5:
+    elif len(args) >= 5 and isinstance(args[4], dict):
         args[4].update(context)
     else:
         kwargs["headers"] = context
@@ -86,9 +100,9 @@ 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:
+    if "headers" in kwargs and isinstance(kwargs["headers"], dict):
         kwargs["headers"].update(context)
-    elif len(args) >= 4:
+    elif len(args) >= 4 and isinstance(args[3], dict):
         args[3].update(context)
     else:
         kwargs["headers"] = context

tests/integration/handle.py

@@ -1,5 +1,3 @@
-import json
-
 from datadog_lambda.metric import lambda_metric
 from datadog_lambda.wrapper import datadog_lambda_wrapper
 

tests/integration/http_requests.py

@@ -1,4 +1,3 @@
-import json
 import requests
 
 from datadog_lambda.metric import lambda_metric
@@ -12,7 +11,7 @@ def handle(event, context):
         "tests.integration.count", 21, tags=["test:integration", "role:hello"]
     )
 
-    us_response = requests.get("https://ip-ranges.datadoghq.com/")
-    eu_response = requests.get("https://ip-ranges.datadoghq.eu/")
+    requests.get("https://ip-ranges.datadoghq.com/")
+    requests.get("https://ip-ranges.datadoghq.eu/")
 
     return {"statusCode": 200, "body": {"message": "hello, dog!"}}

tests/test_patch.py

@@ -7,20 +7,20 @@
 
 from datadog_lambda.patch import (
     _patch_httplib,
-    _patch_requests,
+    _ensure_patch_requests,
 )
 from datadog_lambda.constants import TraceHeader
 
 
 class TestPatchHTTPClients(unittest.TestCase):
 
     def setUp(self):
-        patcher = patch('datadog_lambda.patch.get_dd_trace_context')
+        patcher = patch("datadog_lambda.patch.get_dd_trace_context")
         self.mock_get_dd_trace_context = patcher.start()
         self.mock_get_dd_trace_context.return_value = {
-            TraceHeader.TRACE_ID: '123',
-            TraceHeader.PARENT_ID: '321',
-            TraceHeader.SAMPLING_PRIORITY: '2',
+            TraceHeader.TRACE_ID: "123",
+            TraceHeader.PARENT_ID: "321",
+            TraceHeader.SAMPLING_PRIORITY: "2",
         }
         self.addCleanup(patcher.stop)
 
@@ -34,10 +34,29 @@ def test_patch_httplib(self):
         self.mock_get_dd_trace_context.assert_called()
 
     def test_patch_requests(self):
-        _patch_requests()
+        _ensure_patch_requests()
         import requests
         r = requests.get("https://www.datadoghq.com/")
         self.mock_get_dd_trace_context.assert_called()
-        self.assertEqual(r.request.headers[TraceHeader.TRACE_ID], '123')
-        self.assertEqual(r.request.headers[TraceHeader.PARENT_ID], '321')
-        self.assertEqual(r.request.headers[TraceHeader.SAMPLING_PRIORITY], '2')
+        self.assertEqual(r.request.headers[TraceHeader.TRACE_ID], "123")
+        self.assertEqual(r.request.headers[TraceHeader.PARENT_ID], "321")
+        self.assertEqual(r.request.headers[TraceHeader.SAMPLING_PRIORITY], "2")
+
+    def test_patch_requests_with_headers(self):
+        _ensure_patch_requests()
+        import requests
+        r = requests.get("https://www.datadoghq.com/", headers={"key": "value"})
+        self.mock_get_dd_trace_context.assert_called()
+        self.assertEqual(r.request.headers["key"], "value")
+        self.assertEqual(r.request.headers[TraceHeader.TRACE_ID], "123")
+        self.assertEqual(r.request.headers[TraceHeader.PARENT_ID], "321")
+        self.assertEqual(r.request.headers[TraceHeader.SAMPLING_PRIORITY], "2")
+
+    def test_patch_requests_with_headers_none(self):
+        _ensure_patch_requests()
+        import requests
+        r = requests.get("https://www.datadoghq.com/", headers=None)
+        self.mock_get_dd_trace_context.assert_called()
+        self.assertEqual(r.request.headers[TraceHeader.TRACE_ID], "123")
+        self.assertEqual(r.request.headers[TraceHeader.PARENT_ID], "321")
+        self.assertEqual(r.request.headers[TraceHeader.SAMPLING_PRIORITY], "2")