Add propagator injection for botocore

This commit is contained in:
Nathaniel Ruiz Nowell
2020-11-12 14:03:37 -08:00
parent cb70679ae8
commit 58f77f5d7e
2 changed files with 68 additions and 0 deletions

View File

@ -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)

View File

@ -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)