Add propagator injection for botocore

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py

@@ -59,6 +59,7 @@ from wrapt import ObjectProxy, wrap_function_wrapper
 from opentelemetry.instrumentation.botocore.version import __version__
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.sdk.trace import Resource
+from opentelemetry import propagators
 from opentelemetry.trace import SpanKind, get_tracer
 
 logger = logging.getLogger(__name__)
@@ -85,6 +86,12 @@ class BotocoreInstrumentor(BaseInstrumentor):
             self._patched_api_call,
         )
 
+        wrap_function_wrapper(
+            "botocore.endpoint",
+            "Endpoint.prepare_request",
+            self._patched_prepare_request,
+        )
+
     def _uninstrument(self, **kwargs):
         unwrap(BaseClient, "_make_api_call")
 
@@ -144,6 +151,12 @@ class BotocoreInstrumentor(BaseInstrumentor):
 
             return result
 
+    def _patched_prepare_request(self, wrapped, instance, args, kwargs):
+        request = args[0]
+        headers = request.headers
+        propagators.inject(type(headers).__setitem__, headers)
+        return wrapped(*args, **kwargs)
+
 
 def unwrap(obj, attr):
     function = getattr(obj, attr, None)

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py

@@ -26,7 +26,9 @@ from moto import (  # pylint: disable=import-error
 )
 
 from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
+from opentelemetry import propagators
 from opentelemetry.sdk.resources import Resource
+from opentelemetry.test.mock_textmap import MockTextMapPropagator
 from opentelemetry.test.test_base import TestBase
 
 
@@ -275,3 +277,56 @@ class TestBotocoreInstrumentor(TestBase):
 
         # checking for protection on sts against security leak
         self.assertTrue("params" not in span.attributes.keys())
+
+    @mock_ec2
+    def test_propagator_injects_into_request(self):
+        headers = {}
+        previous_propagator = propagators.get_global_textmap()
+
+        def check_headers(**kwargs):
+            nonlocal headers
+            headers = kwargs["request"].headers
+
+        try:
+            propagators.set_global_textmap(MockTextMapPropagator())
+
+            ec2 = self.session.create_client("ec2", region_name="us-west-2")
+            ec2.meta.events.register_first(
+                "before-send.ec2.DescribeInstances", check_headers
+            )
+            ec2.describe_instances()
+
+            spans = self.memory_exporter.get_finished_spans()
+            assert spans
+            span = spans[0]
+            self.assertEqual(len(spans), 1)
+            self.assertEqual(span.attributes["aws.agent"], "botocore")
+            self.assertEqual(span.attributes["aws.region"], "us-west-2")
+            self.assertEqual(
+                span.attributes["aws.operation"], "DescribeInstances"
+            )
+            assert_span_http_status_code(span, 200)
+            self.assertEqual(
+                span.resource,
+                Resource(
+                    attributes={
+                        "endpoint": "ec2",
+                        "operation": "describeinstances",
+                    }
+                ),
+            )
+            self.assertEqual(span.name, "ec2.command")
+
+            self.assertIn(MockTextMapPropagator.TRACE_ID_KEY, headers)
+            self.assertEqual(
+                str(span.get_span_context().trace_id),
+                headers[MockTextMapPropagator.TRACE_ID_KEY],
+            )
+            self.assertIn(MockTextMapPropagator.SPAN_ID_KEY, headers)
+            self.assertEqual(
+                str(span.get_span_context().span_id),
+                headers[MockTextMapPropagator.SPAN_ID_KEY],
+            )
+
+        finally:
+            propagators.set_global_textmap(previous_propagator)