fix(botocore): inject trace context into stepfunction start_execution…

… input (#10262)

In #7514 , I added a feature to patch start_execution calls to AWS
stepfunctions and inject trace context into the input object. That
feature suffered a regression in #8937 , and the whole stepfunctions
integration was refactored to use the Core API in #9005 . This PR
re-adds the injection of trace context into the start_execution input,
and adds a few unit tests to detect the same regression in the future.


## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

ddtrace/_trace/trace_handlers.py

@@ -17,6 +17,7 @@
 from ddtrace.constants import SPAN_KIND
 from ddtrace.constants import SPAN_MEASURED_KEY
 from ddtrace.contrib import trace_utils
+from ddtrace.contrib.internal.botocore.constants import BOTOCORE_STEPFUNCTIONS_INPUT_KEY
 from ddtrace.contrib.trace_utils import _get_request_header_user_agent
 from ddtrace.contrib.trace_utils import _set_url_tag
 from ddtrace.ext import SpanKind
@@ -651,6 +652,12 @@ def _on_botocore_update_messages(ctx, span, _, trace_data, __, message=None):
     HTTPPropagator.inject(context, trace_data)
 
 
+def _on_botocore_patched_stepfunctions_update_input(ctx, span, _, trace_data, __):
+    context = span.context if span else ctx[ctx["call_key"]].context
+    HTTPPropagator.inject(context, trace_data["_datadog"])
+    ctx.set_item(BOTOCORE_STEPFUNCTIONS_INPUT_KEY, trace_data)
+
+
 def _on_botocore_patched_bedrock_api_call_started(ctx, request_params):
     span = ctx[ctx["call_key"]]
     integration = ctx["bedrock_integration"]
@@ -790,7 +797,7 @@ def listen():
     core.on("botocore.sqs_sns.update_messages", _on_botocore_update_messages)
     core.on("botocore.patched_stepfunctions_api_call.started", _on_botocore_patched_api_call_started)
     core.on("botocore.patched_stepfunctions_api_call.exception", _on_botocore_patched_api_call_exception)
-    core.on("botocore.stepfunctions.update_messages", _on_botocore_update_messages)
+    core.on("botocore.stepfunctions.update_input", _on_botocore_patched_stepfunctions_update_input)
     core.on("botocore.eventbridge.update_messages", _on_botocore_update_messages)
     core.on("botocore.client_context.update_messages", _on_botocore_update_messages)
     core.on("botocore.patched_bedrock_api_call.started", _on_botocore_patched_bedrock_api_call_started)

ddtrace/contrib/internal/botocore/constants.py

@@ -0,0 +1 @@
+BOTOCORE_STEPFUNCTIONS_INPUT_KEY = "botocore_stepfunctions_input"

ddtrace/contrib/internal/botocore/services/stepfunctions.py

@@ -5,6 +5,7 @@
 import botocore.exceptions
 
 from ddtrace import config
+from ddtrace.contrib.internal.botocore.constants import BOTOCORE_STEPFUNCTIONS_INPUT_KEY
 from ddtrace.contrib.trace_utils import ext_service
 from ddtrace.ext import SpanTypes
 from ddtrace.internal import core
@@ -34,8 +35,11 @@ def update_stepfunction_input(ctx: core.ExecutionContext, params: Any) -> None:
         return
 
     input_obj["_datadog"] = {}
-
-    core.dispatch("botocore.stepfunctions.update_input", [ctx, None, None, input_obj["_datadog"], None])
+    core.dispatch("botocore.stepfunctions.update_input", [ctx, None, None, input_obj, None])
+    updated_input_obj = ctx.get_item(BOTOCORE_STEPFUNCTIONS_INPUT_KEY)
+    if updated_input_obj:
+        input_json_str = json.dumps(updated_input_obj)
+        params["input"] = input_json_str
 
 
 def patched_stepfunction_api_call(original_func, instance, args, kwargs: Dict, function_vars: Dict):

releasenotes/notes/fix-botocore-stepfunction-patching-9993630697956278.yaml

@@ -0,0 +1,4 @@
+fixes:
+  - |
+    botocore: This fix resolves a regression where trace context was not being injected into the input of Stepfunction
+      start_execution commands. This re-enables distributed tracing when a Python service invokes a properly instrumented Step Function.

tests/contrib/botocore/test.py

@@ -1046,9 +1046,11 @@ def test_stepfunctions_send_start_execution_trace_injection(self):
             roleArn="arn:aws:iam::012345678901:role/DummyRole",
         )
         Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf)
-        sf.start_execution(
-            stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:lincoln", input='{"baz":1}'
-        )
+        start_execution_dict = {
+            "stateMachineArn": "arn:aws:states:us-west-2:000000000000:stateMachine:lincoln",
+            "input": '{"baz": 1}',
+        }
+        sf.start_execution(**start_execution_dict)
         # I've tried to find a way to make Moto show me the input to the execution, but can't get that to work.
         spans = self.get_spans()
         assert spans

tests/contrib/botocore/test_stepfunctions.py

@@ -0,0 +1,99 @@
+import json
+
+from ddtrace import Pin
+from ddtrace.contrib.internal.botocore.services.stepfunctions import update_stepfunction_input
+from ddtrace.ext import SpanTypes
+from ddtrace.internal import core
+
+
+def test_update_stepfunction_input():
+    params = {
+        "stateMachineArn": "arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine",
+        "input": "{}",
+    }
+    pin = Pin()
+    with core.context_with_data(
+        "botocore.patched_stepfunctions_api_call",
+        span_name="states.command",
+        service="aws.states",
+        span_type=SpanTypes.HTTP,
+        call_key="patched_stepfunctions_api_call",
+        instance=None,
+        args=(
+            "StartExecution",
+            {
+                "stateMachineArn": "arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine",
+                "input": "{}",
+            },
+        ),
+        params=params,
+        endpoint_name="states",
+        operation="StartExecution",
+        pin=pin,
+    ) as ctx:
+        update_stepfunction_input(ctx, params)
+        assert params["input"]
+        input_obj = json.loads(params["input"])
+        assert "_datadog" in input_obj
+        assert "x-datadog-trace-id" in input_obj["_datadog"]
+        assert "x-datadog-parent-id" in input_obj["_datadog"]
+        assert "x-datadog-tags" in input_obj["_datadog"]
+        assert "_dd.p.tid" in input_obj["_datadog"]["x-datadog-tags"]
+
+
+def test_update_stepfunction_input_does_not_mutate_non_dict_input():
+    params = {
+        "stateMachineArn": "arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine",
+        "input": "hello",
+    }
+    pin = Pin()
+    with core.context_with_data(
+        "botocore.patched_stepfunctions_api_call",
+        span_name="states.command",
+        service="aws.states",
+        span_type=SpanTypes.HTTP,
+        call_key="patched_stepfunctions_api_call",
+        instance=None,
+        args=(
+            "StartExecution",
+            {
+                "stateMachineArn": "arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine",
+                "input": "hello",
+            },
+        ),
+        params=params,
+        endpoint_name="states",
+        operation="StartExecution",
+        pin=pin,
+    ) as ctx:
+        update_stepfunction_input(ctx, params)
+        assert params["input"]
+        assert params["input"] == "hello"
+
+    params = {
+        "stateMachineArn": "arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine",
+        "input": "[1, 2, 3]",
+    }
+    pin = Pin()
+    with core.context_with_data(
+        "botocore.patched_stepfunctions_api_call",
+        span_name="states.command",
+        service="aws.states",
+        span_type=SpanTypes.HTTP,
+        call_key="patched_stepfunctions_api_call",
+        instance=None,
+        args=(
+            "StartExecution",
+            {
+                "stateMachineArn": "arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine",
+                "input": "[1, 2, 3]",
+            },
+        ),
+        params=params,
+        endpoint_name="states",
+        operation="StartExecution",
+        pin=pin,
+    ) as ctx:
+        update_stepfunction_input(ctx, params)
+        assert params["input"]
+        assert params["input"] == "[1, 2, 3]"