update botocore patch function signatures

Author: yiyuan-he

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py

@@ -31,6 +31,7 @@
     _AttributeMapT,
     _AwsSdkCallContext,
     _AwsSdkExtension,
+    _BotocoreInstrumentorContext,
     _BotoResultT,
 )
 from opentelemetry.trace.span import Span
@@ -192,7 +193,7 @@ def extract_attributes(self, attributes: _AttributeMapT):
             if request_param_value:
                 attributes[attribute_key] = request_param_value
 
-    def on_success(self, span: Span, result: _BotoResultT):
+    def on_success(self, span: Span, result: _BotoResultT, instrumentor_context: _BotocoreInstrumentorContext):
         if self._operation_class is None:
             return
 
@@ -229,7 +230,7 @@ class _BedrockExtension(_AwsSdkExtension):
     """
 
     # pylint: disable=no-self-use
-    def on_success(self, span: Span, result: _BotoResultT):
+    def on_success(self, span: Span, result: _BotoResultT, instrumentor_context: _BotocoreInstrumentorContext):
         # _GUARDRAIL_ID can only be retrieved from the response, not from the request
         guardrail_id = result.get(_GUARDRAIL_ID)
         if guardrail_id:
@@ -333,7 +334,7 @@ def _set_if_not_none(attributes, key, value):
             attributes[key] = value
 
     # pylint: disable=too-many-branches
-    def on_success(self, span: Span, result: Dict[str, Any]):
+    def on_success(self, span: Span, result: Dict[str, Any], instrumentor_context: _BotocoreInstrumentorContext):
         model_id = self._call_context.params.get(_MODEL_ID)
 
         if not model_id:

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py

@@ -25,7 +25,12 @@
 from opentelemetry.instrumentation.botocore.extensions.lmbd import _LambdaExtension
 from opentelemetry.instrumentation.botocore.extensions.sns import _SnsExtension
 from opentelemetry.instrumentation.botocore.extensions.sqs import _SqsExtension
-from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension, _BotoResultT
+from opentelemetry.instrumentation.botocore.extensions.types import (
+    _AttributeMapT,
+    _AwsSdkExtension,
+    _BotocoreInstrumentorContext,
+    _BotoResultT,
+)
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace.span import Span
 
@@ -75,8 +80,8 @@ def patch_extract_attributes(self, attributes: _AttributeMapT):
 
     old_on_success = _LambdaExtension.on_success
 
-    def patch_on_success(self, span: Span, result: _BotoResultT):
-        old_on_success(self, span, result)
+    def patch_on_success(self, span: Span, result: _BotoResultT, instrumentor_context: _BotocoreInstrumentorContext):
+        old_on_success(self, span, result, instrumentor_context)
         lambda_configuration = result.get("Configuration", {})
         function_arn = lambda_configuration.get("FunctionArn")
         if function_arn:
@@ -180,8 +185,8 @@ def patch_extract_attributes(self, attributes: _AttributeMapT):
 
     old_on_success = _SqsExtension.on_success
 
-    def patch_on_success(self, span: Span, result: _BotoResultT):
-        old_on_success(self, span, result)
+    def patch_on_success(self, span: Span, result: _BotoResultT, instrumentor_context: _BotocoreInstrumentorContext):
+        old_on_success(self, span, result, instrumentor_context)
         queue_url = result.get("QueueUrl")
         if queue_url:
             span.set_attribute(AWS_SQS_QUEUE_URL, queue_url)
@@ -243,7 +248,7 @@ def extract_attributes(self, attributes: _AttributeMapT):
             attributes[AWS_SECRETSMANAGER_SECRET_ARN] = secret_id
 
     # pylint: disable=no-self-use
-    def on_success(self, span: Span, result: _BotoResultT):
+    def on_success(self, span: Span, result: _BotoResultT, instrumentor_context: _BotocoreInstrumentorContext):
         secret_arn = result.get("ARN")
         if secret_arn:
             span.set_attribute(AWS_SECRETSMANAGER_SECRET_ARN, secret_arn)

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py

@@ -678,6 +678,7 @@ def _do_on_success(
 ) -> Dict[str, str]:
     span_mock: Span = MagicMock()
     mock_call_context = MagicMock()
+    mock_instrumentor_context = MagicMock()
     span_attributes: Dict[str, str] = {}
 
     def set_side_effect(set_key, set_value):
@@ -692,6 +693,6 @@ def set_side_effect(set_key, set_value):
         mock_call_context.params = params
 
     extension = _KNOWN_EXTENSIONS[service_name]()(mock_call_context)
-    extension.on_success(span_mock, result)
+    extension.on_success(span_mock, result, mock_instrumentor_context)
 
     return span_attributes