Update trace propagation format for botocore direct invocations (#2639)

* Update trace propagation format for botocore direct invocations

* Update client context format in one other place

* Add config for add_datadog_object

* Refactor modify_client_context and inject_trace_to_client_context

* black

* default false

* Add a test for the config flag

* Document the config variable

* Add release notes

* add a bunch of newlines for flake8

* Defaut -> Default

* add js and clientContext to wordlist

* Update ddtrace/contrib/botocore/patch.py

Mutate client_context_object instead of returning

Co-authored-by: Kyle Verhoog <[email protected]>

* Update ddtrace/contrib/botocore/patch.py

Co-authored-by: Kyle Verhoog <[email protected]>

* Trimmed down on excess documentation

* rename config value to BOTOCORE_INVOKE_WITH_LEGACY_CONTEXT

* run black

* Remove clientContext form wordlist

* Improve documentation

* upgrades -> upgrade

* re-write release notes

* fix spelling of fnuctions

Co-authored-by: Kyle Verhoog <[email protected]>
Co-authored-by: Brett Langdon <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 4 people

ddtrace/contrib/botocore/__init__.py

@@ -24,6 +24,18 @@
 
    Default: ``True``
 
+.. py:data:: ddtrace.config.botocore['invoke_with_legacy_context']
+
+    This preserves legacy behavior when tracing directly invoked Python and Node Lambda
+    functions instrumented with datadog-lambda-python < v41 or datadog-lambda-js < v3.58.0.
+
+    Legacy support for older libraries is available with
+    ``ddtrace.config.botocore.invoke_with_legacy_context = True`` or by setting the environment
+    variable ``DD_BOTOCORE_INVOKE_WITH_LEGACY_CONTEXT=true``.
+
+
+    Default: ``False``
+
 
 Example::
 

ddtrace/contrib/botocore/patch.py

@@ -37,6 +37,7 @@
     "botocore",
     {
         "distributed_tracing": get_env("botocore", "distributed_tracing", default=True),
+        "invoke_with_legacy_context": get_env("botocore", "invoke_with_legacy_context", default=False),
     },
 )
 
@@ -68,36 +69,35 @@ def inject_trace_to_sqs_message(args, span):
     inject_trace_data_to_message_attributes(trace_data, params)
 
 
-def modify_client_context(client_context_base64, trace_headers):
-    try:
-        client_context_json = base64.b64decode(client_context_base64).decode("utf-8")
-        client_context_object = json.loads(client_context_json)
-
-        if "custom" in client_context_object:
-            client_context_object["custom"]["_datadog"] = trace_headers
-        else:
-            client_context_object["custom"] = {"_datadog": trace_headers}
+def modify_client_context(client_context_object, trace_headers):
+    if config.botocore["invoke_with_legacy_context"]:
+        trace_headers = {"_datadog": trace_headers}
 
-        new_context = base64.b64encode(json.dumps(client_context_object).encode("utf-8")).decode("utf-8")
-        return new_context
-    except Exception:
-        log.warning("malformed client_context=%s", client_context_base64, exc_info=True)
-        return client_context_base64
+    if "custom" in client_context_object:
+        client_context_object["custom"].update(trace_headers)
+    else:
+        client_context_object["custom"] = trace_headers
 
 
 def inject_trace_to_client_context(args, span):
     trace_headers = {}
     HTTPPropagator.inject(span.context, trace_headers)
-
+    client_context_object = {}
     params = args[1]
     if "ClientContext" in params:
-        params["ClientContext"] = modify_client_context(params["ClientContext"], trace_headers)
-    else:
-        trace_headers = {}
-        HTTPPropagator.inject(span.context, trace_headers)
-        client_context_object = {"custom": {"_datadog": trace_headers}}
+        try:
+            client_context_json = base64.b64decode(params["ClientContext"]).decode("utf-8")
+            client_context_object = json.loads(client_context_json)
+        except Exception:
+            log.warning("malformed client_context=%s", params["ClientContext"], exc_info=True)
+            return
+    modify_client_context(client_context_object, trace_headers)
+    try:
         json_context = json.dumps(client_context_object).encode("utf-8")
-        params["ClientContext"] = base64.b64encode(json_context).decode("utf-8")
+    except Exception:
+        log.warning("unable to encode modified client context as json: %s", client_context_object, exc_info=True)
+        return
+    params["ClientContext"] = base64.b64encode(json_context).decode("utf-8")
 
 
 def patch():

docs/spelling_wordlist.txt

@@ -66,6 +66,7 @@ integration
 integrations
 JSON
 jinja
+js
 kombu
 kubernetes
 kwarg

releasenotes/notes/botocore-46f5bf04f64777b4.yaml

@@ -0,0 +1,8 @@
+---
+upgrade:
+  - |
+    botocore: Update trace propagation format for directly invoked Lambda functions.
+    This breaks compatibility with Lambda functions instrumented with datadog-lambda-python < v41 or datadog-lambda-js < v3.57.0.
+    Please upgrade datadog-lambda-* in invoked lambda functions, or engage legacy compatibility mode in one of two ways:
+    - ddtrace.config.botocore.invoke_with_legacy_context = True
+    - DD_BOTOCORE_INVOKE_WITH_LEGACY_CONTEXT=true

tests/contrib/botocore/test.py

@@ -526,11 +526,55 @@ def test_lambda_invoke_no_context_client(self):
         context_json = base64.b64decode(context_b64.encode()).decode()
         context_obj = json.loads(context_json)
 
-        self.assertEqual(context_obj["custom"]["_datadog"][HTTP_HEADER_TRACE_ID], str(span.trace_id))
-        self.assertEqual(context_obj["custom"]["_datadog"][HTTP_HEADER_PARENT_ID], str(span.span_id))
+        self.assertEqual(context_obj["custom"][HTTP_HEADER_TRACE_ID], str(span.trace_id))
+        self.assertEqual(context_obj["custom"][HTTP_HEADER_PARENT_ID], str(span.span_id))
 
         lamb.delete_function(FunctionName="ironmaiden")
 
+    @mock_lambda
+    def test_lambda_invoke_with_old_style_trace_propagation(self):
+        with self.override_config("botocore", dict(invoke_with_legacy_context=True)):
+            lamb = self.session.create_client("lambda", region_name="us-west-2", endpoint_url="http://localhost:4566")
+            lamb.create_function(
+                FunctionName="ironmaiden",
+                Runtime="python3.7",
+                Role="test-iam-role",
+                Handler="lambda_function.lambda_handler",
+                Code={
+                    "ZipFile": get_zip_lambda(),
+                },
+                Publish=True,
+                Timeout=30,
+                MemorySize=128,
+            )
+
+            Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(lamb)
+
+            lamb.invoke(
+                FunctionName="ironmaiden",
+                Payload=json.dumps({}),
+            )
+
+            spans = self.get_spans()
+            assert spans
+            span = spans[0]
+
+            self.assertEqual(len(spans), 1)
+            self.assertEqual(span.get_tag("aws.region"), "us-west-2")
+            self.assertEqual(span.get_tag("aws.operation"), "Invoke")
+            assert_is_measured(span)
+            assert_span_http_status_code(span, 200)
+            self.assertEqual(span.service, "test-botocore-tracing.lambda")
+            self.assertEqual(span.resource, "lambda.invoke")
+            context_b64 = span.get_tag("params.ClientContext")
+            context_json = base64.b64decode(context_b64.encode()).decode()
+            context_obj = json.loads(context_json)
+
+            self.assertEqual(context_obj["custom"]["_datadog"][HTTP_HEADER_TRACE_ID], str(span.trace_id))
+            self.assertEqual(context_obj["custom"]["_datadog"][HTTP_HEADER_PARENT_ID], str(span.span_id))
+
+            lamb.delete_function(FunctionName="ironmaiden")
+
     @mock_lambda
     def test_lambda_invoke_distributed_tracing_off(self):
         with self.override_config("botocore", dict(distributed_tracing=False)):
@@ -610,8 +654,8 @@ def test_lambda_invoke_with_context_client(self):
         context_obj = json.loads(context_json)
 
         self.assertEqual(context_obj["custom"]["foo"], "bar")
-        self.assertEqual(context_obj["custom"]["_datadog"][HTTP_HEADER_TRACE_ID], str(span.trace_id))
-        self.assertEqual(context_obj["custom"]["_datadog"][HTTP_HEADER_PARENT_ID], str(span.span_id))
+        self.assertEqual(context_obj["custom"][HTTP_HEADER_TRACE_ID], str(span.trace_id))
+        self.assertEqual(context_obj["custom"][HTTP_HEADER_PARENT_ID], str(span.span_id))
 
         lamb.delete_function(FunctionName="megadeth")