mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-07-31 14:11:50 +08:00
Add propagator injection for botocore
This commit is contained in:
@ -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)
|
||||
|
@ -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)
|
||||
|
Reference in New Issue
Block a user