@@ -101,22 +101,14 @@ def response_hook(span, service_name, operation_name, result):
101101 _SUPPRESS_INSTRUMENTATION_KEY ,
102102 unwrap ,
103103)
104- from opentelemetry .propagate import inject
104+ from opentelemetry .propagators . aws . aws_xray_propagator import AwsXRayPropagator
105105from opentelemetry .semconv .trace import SpanAttributes
106106from opentelemetry .trace import get_tracer
107107from opentelemetry .trace .span import Span
108108
109109logger = logging .getLogger (__name__ )
110110
111111
112- # pylint: disable=unused-argument
113- def _patched_endpoint_prepare_request (wrapped , instance , args , kwargs ):
114- request = args [0 ]
115- headers = request .headers
116- inject (headers )
117- return wrapped (* args , ** kwargs )
118-
119-
120112class BotocoreInstrumentor (BaseInstrumentor ):
121113 """An instrumentor for Botocore.
122114
@@ -127,6 +119,7 @@ def __init__(self):
127119 super ().__init__ ()
128120 self .request_hook = None
129121 self .response_hook = None
122+ self .propagator = AwsXRayPropagator ()
130123
131124 def instrumentation_dependencies (self ) -> Collection [str ]:
132125 return _instruments
@@ -140,6 +133,10 @@ def _instrument(self, **kwargs):
140133 self .request_hook = kwargs .get ("request_hook" )
141134 self .response_hook = kwargs .get ("response_hook" )
142135
136+ propagator = kwargs .get ("propagator" )
137+ if propagator is not None :
138+ self .propagator = propagator
139+
143140 wrap_function_wrapper (
144141 "botocore.client" ,
145142 "BaseClient._make_api_call" ,
@@ -149,13 +146,26 @@ def _instrument(self, **kwargs):
149146 wrap_function_wrapper (
150147 "botocore.endpoint" ,
151148 "Endpoint.prepare_request" ,
152- _patched_endpoint_prepare_request ,
149+ self . _patched_endpoint_prepare_request ,
153150 )
154151
155152 def _uninstrument (self , ** kwargs ):
156153 unwrap (BaseClient , "_make_api_call" )
157154 unwrap (Endpoint , "prepare_request" )
158155
156+ # pylint: disable=unused-argument
157+ def _patched_endpoint_prepare_request (
158+ self , wrapped , instance , args , kwargs
159+ ):
160+ request = args [0 ]
161+ headers = request .headers
162+
163+ # Only the x-ray header is propagated by AWS services. Using any
164+ # other propagator will lose the trace context.
165+ self .propagator .inject (headers )
166+
167+ return wrapped (* args , ** kwargs )
168+
159169 # pylint: disable=too-many-branches
160170 def _patched_api_call (self , original_func , instance , args , kwargs ):
161171 if context_api .get_value (_SUPPRESS_INSTRUMENTATION_KEY ):
0 commit comments