Merge pull request #181 from NathanielRN/botocore-propagation-injection

Add propagator injection for botocore calls

instrumentation/opentelemetry-instrumentation-botocore/CHANGELOG.md

@@ -1,6 +1,8 @@
 # Changelog
 
 ## Unreleased
+- Add propagator injection for botocore calls
+  ([#181](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/181))
 
 ## Version 0.13b0
 

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

@@ -56,6 +56,7 @@ import logging
 from botocore.client import BaseClient
 from wrapt import ObjectProxy, wrap_function_wrapper
 
+from opentelemetry import propagators
 from opentelemetry.instrumentation.botocore.version import __version__
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.sdk.trace import Resource
@@ -64,6 +65,14 @@ from opentelemetry.trace import SpanKind, get_tracer
 logger = logging.getLogger(__name__)
 
 
+# pylint: disable=unused-argument
+def _patched_endpoint_prepare_request(wrapped, instance, args, kwargs):
+    request = args[0]
+    headers = request.headers
+    propagators.inject(type(headers).__setitem__, headers)
+    return wrapped(*args, **kwargs)
+
+
 class BotocoreInstrumentor(BaseInstrumentor):
     """A instrumentor for Botocore
 
@@ -85,6 +94,12 @@ class BotocoreInstrumentor(BaseInstrumentor):
             self._patched_api_call,
         )
 
+        wrap_function_wrapper(
+            "botocore.endpoint",
+            "Endpoint.prepare_request",
+            _patched_endpoint_prepare_request,
+        )
+
     def _uninstrument(self, **kwargs):
         unwrap(BaseClient, "_make_api_call")
 

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

@@ -25,8 +25,10 @@ from moto import (  # pylint: disable=import-error
     mock_sqs,
 )
 
+from opentelemetry import propagators
 from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
 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)