Try never break the wrapped lambda function

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

datadog_lambda/wrapper.py

@@ -53,37 +53,45 @@ class _LambdaDecorator(object):
  and extracts/injects trace context.
  """
 
- _force_new = False
+ _force_wrap = False
 
  def __new__(cls, func):
  """
  If the decorator is accidentally applied to the same function multiple times,
- only the first one takes effect.
+ wrap only once.
 
- If _force_new, always return a real decorator, useful for unit tests.
+ If _force_wrap, always return a real decorator, useful for unit tests.
  """
- if cls._force_new or not getattr(func, "_dd_wrapped", False):
- wrapped = super(_LambdaDecorator, cls).__new__(cls)
- wrapped._dd_wrapped = True
- return wrapped
- else:
- return _NoopDecorator(func)
+ try:
+ if cls._force_wrap or not isinstance(func, _LambdaDecorator):
+ wrapped = super(_LambdaDecorator, cls).__new__(cls)
+ logger.debug("datadog_lambda_wrapper wrapped")
+ return wrapped
+ else:
+ logger.debug("datadog_lambda_wrapper already wrapped")
+ return _NoopDecorator(func)
+ except Exception:
+ traceback.print_exc()
+ return func
 
  def __init__(self, func):
  """Executes when the wrapped function gets wrapped"""
- self.func = func
- self.flush_to_log = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
- self.logs_injection = (
- os.environ.get("DD_LOGS_INJECTION", "true").lower() == "true"
- )
-
- # Inject trace correlation ids to logs
- if self.logs_injection:
- inject_correlation_ids()
-
- # Patch HTTP clients to propagate Datadog trace context
- patch_all()
- logger.debug("datadog_lambda_wrapper initialized")
+ try:
+ self.func = func
+ self.flush_to_log = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
+ self.logs_injection = (
+ os.environ.get("DD_LOGS_INJECTION", "true").lower() == "true"
+ )
+
+ # Inject trace correlation ids to logs
+ if self.logs_injection:
+ inject_correlation_ids()
+
+ # Patch HTTP clients to propagate Datadog trace context
+ patch_all()
+ logger.debug("datadog_lambda_wrapper initialized")
+ except Exception:
+ traceback.print_exc()
 
  def __call__(self, event, context, **kwargs):
  """Executes when the wrapped function gets called"""
@@ -97,21 +105,23 @@ def __call__(self, event, context, **kwargs):
  self._after(event, context)
 
  def _before(self, event, context):
- set_cold_start()
  try:
+ set_cold_start()
  submit_invocations_metric(context)
  # Extract Datadog trace context from incoming requests
  extract_dd_trace_context(event)
 
  # Set log correlation ids using extracted trace context
  set_correlation_ids()
+ logger.debug("datadog_lambda_wrapper _before() done")
  except Exception:
  traceback.print_exc()
 
  def _after(self, event, context):
  try:
  if not self.flush_to_log:
  lambda_stats.flush(float("inf"))
+ logger.debug("datadog_lambda_wrapper _after() done")
  except Exception:
  traceback.print_exc()
 

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")

tests/test_wrapper.py

@@ -26,7 +26,7 @@ class TestDatadogLambdaWrapper(unittest.TestCase):
  def setUp(self):
  # Force @datadog_lambda_wrapper to always create a real
  # (not no-op) wrapper.
- datadog_lambda_wrapper._force_new = True
+ datadog_lambda_wrapper._force_wrap = True
 
  patcher = patch("datadog_lambda.metric.lambda_stats")
  self.mock_metric_lambda_stats = patcher.start()
@@ -265,9 +265,9 @@ def test_only_one_wrapper_in_use(self):
  def lambda_handler(event, context):
  lambda_metric("test.metric", 100)
 
- # Turn off _force_new to emulate the nested wrapper scenario,
+ # Turn off _force_wrap to emulate the nested wrapper scenario,
  # the second @datadog_lambda_wrapper should actually be no-op.
- datadog_lambda_wrapper._force_new = False
+ datadog_lambda_wrapper._force_wrap = False
 
  lambda_handler_double_wrapped = datadog_lambda_wrapper(lambda_handler)