adapt to both before-call and before-sign

Signed-off-by: Varsha GS <[email protected]>

instana/instrumentation/boto3_inst.py

@@ -46,16 +46,16 @@ def lambda_inject_context(payload, scope):
             logger.debug("non-fatal lambda_inject_context: ", exc_info=True)
 
 
-    @wrapt.patch_function_wrapper("botocore.hooks", "HierarchicalEmitter.emit")
-    def emit_request_created_with_instana(wrapped, instance, args, kwargs):
+    @wrapt.patch_function_wrapper("botocore.auth", "SigV4Auth.add_auth")
+    def emit_add_auth_with_instana(wrapped, instance, args, kwargs):
         active_tracer = get_active_tracer()
-        
-        # If we're not tracing or the event emitted is not request-created, just return;
-        if active_tracer is None or args[0].split(".")[0] != "request-created":
+
+        # If we're not tracing, just return;
+        if active_tracer is None:
             return wrapped(*args, **kwargs)
 
         span = active_tracer.active_span
-        extract_custom_headers(span, kwargs["request"].headers)
+        extract_custom_headers(span, args[0].headers)
 
         return wrapped(*args, **kwargs)
 

tests/clients/boto3/test_boto3_s3.py

@@ -296,8 +296,15 @@ def test_request_header_capture(self):
         def add_custom_header_before_call(params, **kwargs):
             params['headers'].update(request_headers)
 
-        # Register the function to an event.
-        event_system.register('before-call.s3.CreateBucket', add_custom_header_before_call)
+        # # Register the function to before-call event.
+        # event_system.register('before-call.s3.CreateBucket', add_custom_header_before_call)
+
+        def _add_header(request, **kwargs):
+            for name, value in request_headers.items():
+                request.headers.add_header(name, value)
+
+        # Register the function to before-sign event.
+        event_system.register_first('before-sign.s3.CreateBucket', _add_header)
 
         with tracer.start_active_span('test'):
             result = self.s3.create_bucket(Bucket="aws_bucket_name")