feat(botocore): patch stepfunctions

ddtrace/contrib/botocore/patch.py

@@ -235,6 +235,46 @@ def inject_trace_to_sqs_or_sns_message(params, span, endpoint_service=None, pin=
     inject_trace_data_to_message_attributes(trace_data, params, endpoint_service)
 
 
+def inject_trace_to_stepfunction_input(params, span):
+    # type: (Any, Span) -> None
+    """
+    :params: contains the params for the current botocore action
+    :span: the span which provides the trace context to be propagated
+
+    Inject the trace headers into the StepFunction input if the input is a JSON string
+    """
+    if "input" not in params:
+        log.warning("Unable to inject context. The StepFunction input had no input.")
+        return
+
+    if params["input"] is None:
+        log.warning("Unable to inject context. The StepFunction input was None.")
+        return
+
+    if isinstance(params["input"], dict):
+        if "_datadog" in params["input"]:
+            log.warning("Input already has trace context.")
+            return
+        params["input"]["_datadog"] = {}
+        HTTPPropagator.inject(span.context, params["input"]["_datadog"])
+        return
+
+    if isinstance(params["input"], str):
+        try:
+            input_obj = json.loads(params["input"])
+        except ValueError:
+            log.warning("Input is not a valid JSON string")
+            return
+
+        input_obj["_datadog"] = {}
+        HTTPPropagator.inject(span.context, input_obj["_datadog"])
+        input_json = json.dumps(input_obj)
+
+        params["input"] = input_json
+
+    log.warning("Unable to inject context. The StepFunction input was not a dict or a JSON string.")
+
+
 def inject_trace_to_eventbridge_detail(params, span):
     # type: (Any, Span) -> None
     """
@@ -574,6 +614,16 @@ def patched_api_call(original_func, instance, args, kwargs):
                         span.name = schematize_cloud_messaging_operation(
                             trace_operation, cloud_provider="aws", cloud_service="sns", direction=SpanDirection.OUTBOUND
                         )
+                    if endpoint_name == "stepfunctions" and (
+                        operation == "start-execution" or operation == "start-sync-execution"
+                    ):
+                        inject_trace_to_stepfunction_input(params, span)
+                        span.name = schematize_cloud_messaging_operation(
+                            trace_operation,
+                            cloud_provider="aws",
+                            cloud_service="stepfunctions",
+                            direction=SpanDirection.OUTBOUND,
+                        )
                 except Exception:
                     log.warning("Unable to inject trace context", exc_info=True)
 

releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    botocore: Add the ability to inject trace context into the input field of botocore stepfunction start_execution and 
+    start_sync_execution calls.

tests/contrib/botocore/test.py

@@ -18,6 +18,7 @@
 from moto import mock_s3
 from moto import mock_sns
 from moto import mock_sqs
+from moto import mock_stepfunctions
 import pytest
 
 from tests.utils import get_128_bit_trace_id_from_headers
@@ -42,6 +43,7 @@
 from ddtrace.internal.schema import DEFAULT_SPAN_SERVICE_NAME
 from ddtrace.internal.utils.version import parse_version
 from ddtrace.propagation.http import HTTP_HEADER_PARENT_ID
+from ddtrace.propagation.http import HTTP_HEADER_TRACE_ID
 from tests.opentracer.utils import init_tracer
 from tests.utils import TracerTestCase
 from tests.utils import assert_is_measured
@@ -885,6 +887,40 @@ def test_schematized_unspecified_service_sqs_client_v1(self):
         assert spans[2].service == DEFAULT_SPAN_SERVICE_NAME
         assert spans[2].name == "aws.sqs.receive"
 
+    @mock_stepfunctions
+    def test_stepfunctions_client(self):
+        sf = self.session.create_client("stepfunctions", region_name="us-west-2")
+        Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf)
+        sf.start_execution(
+            stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo", name="bar", input='{"baz":1}'
+        )
+        spans = self.get_spans()
+        assert spans
+        span = spans[0]
+        assert len(spans) == 1
+        assert span.get_tag("aws.region") == "us-west-2"
+        assert span.get_tag("region") == "us-west-2"
+        assert span.get_tag("aws.operation") == "StartExecution"
+        assert span.get_tag("component") == "botocore"
+        assert span.get_tag("span.kind"), "client"
+        assert_is_measured(span)
+        assert_span_http_status_code(span, 200)
+        assert span.service == "test-botocore-tracing.stepfunctions"
+        assert span.resource == "stepfunctions.startexecution"
+        assert span.get_tag("params.input") == '{"baz":1}'
+
+    def test_stepfunctions_send_start_execution_trace_injection(self):
+        sf = self.session.create_client("stepfunctions", region_name="us-west-2")
+        Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf)
+        sf.start_execution(
+            stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo", name="bar", input='{"baz":1}'
+        )
+        spans = self.get_spans()
+        assert spans
+        span = spans[0]
+        input_obj = json.loads(span.get_tag("params.input"))
+        assert input_obj["_datadog"][HTTP_HEADER_TRACE_ID] == str(span.trace_id)
+
     def _test_kinesis_client(self):
         client = self.session.create_client("kinesis", region_name="us-east-1")
         stream_name = "test"