fix: function signature updates for upstream patches

Author: yiyuan-he

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

@@ -192,7 +192,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=None):
         if self._operation_class is None:
             return
 
@@ -220,6 +220,10 @@ def extract_attributes(self, attributes: _AttributeMapT):
         knowledge_base_id = self._call_context.params.get(_KNOWLEDGE_BASE_ID)
         if knowledge_base_id:
             attributes[AWS_BEDROCK_KNOWLEDGE_BASE_ID] = knowledge_base_id
+            
+    def on_success(self, span: Span, result: _BotoResultT, instrumentor_context=None):
+        # Currently no attributes to extract from the result
+        pass
 
 
 class _BedrockExtension(_AwsSdkExtension):
@@ -229,7 +233,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=None):
         # _GUARDRAIL_ID can only be retrieved from the response, not from the request
         guardrail_id = result.get(_GUARDRAIL_ID)
         if guardrail_id:
@@ -333,7 +337,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=None):
         model_id = self._call_context.params.get(_MODEL_ID)
 
         if not model_id:

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

@@ -75,8 +75,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=None):
+        old_on_success(self, span, result, instrumentor_context)
         lambda_configuration = result.get("Configuration", {})
         function_arn = lambda_configuration.get("FunctionArn")
         if function_arn:
@@ -180,8 +180,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=None):
+        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 +243,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=None):
         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

@@ -147,7 +147,7 @@ def _test_unpatched_botocore_instrumentation(self):
         )
 
         # BedrockRuntime
-        self.assertFalse("bedrock-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a bedrock-runtime extension")
+        self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a bedrock-runtime extension")
 
         # SecretsManager
         self.assertFalse("secretsmanager" in _KNOWN_EXTENSIONS, "Upstream has added a SecretsManager extension")
@@ -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