Add SNS support.

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

@@ -74,8 +74,6 @@ def _apply_botocore_sns_patch() -> None:
     This patch extends the existing upstream extension for SNS. Extensions allow for custom logic for adding
     service-specific information to spans, such as attributes. Specifically, we are adding logic to add
     "aws.sns.topic_arn" attributes to be used to generate AWS_REMOTE_RESOURCE_TYPE and AWS_REMOTE_RESOURCE_IDENTIFIER.
-    Callout that today, the upstream logic adds SpanAttributes.MESSAGING_DESTINATION_NAME,
-    but we are not using it as it can be assigned with TargetArn as well.
     """
     old_extract_attributes = _SnsExtension.extract_attributes
 

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

@@ -1,6 +1,6 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-from typing import Dict
+from typing import Any, Dict
 from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
@@ -9,6 +9,7 @@
 from amazon.opentelemetry.distro.patches._instrumentation_patch import apply_instrumentation_patches
 from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS
 from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace.span import Span
 
 _STREAM_NAME: str = "streamName"
 _BUCKET_NAME: str = "bucketName"
@@ -78,8 +79,6 @@ def _validate_unpatched_botocore_instrumentation(self):
         # SNS
         self.assertTrue("sns" in _KNOWN_EXTENSIONS, "Upstream has removed the SNS extension")
         sns_attributes: Dict[str, str] = _do_extract_sns_attributes()
-        self.assertTrue(SpanAttributes.MESSAGING_SYSTEM in sns_attributes)
-        self.assertEqual(sns_attributes[SpanAttributes.MESSAGING_SYSTEM], "aws.sns")
         self.assertFalse("aws.sns.topic_arn" in sns_attributes)
 
     def _validate_patched_botocore_instrumentation(self):
@@ -109,6 +108,9 @@ def _validate_patched_botocore_instrumentation(self):
         sns_attributes: Dict[str, str] = _do_extract_sns_attributes()
         self.assertTrue("aws.sns.topic_arn" in sns_attributes)
         self.assertEqual(sns_attributes["aws.sns.topic_arn"], _TOPIC_ARN)
+        sns_success_attributes: Dict[str, str] = _do_sns_on_success()
+        self.assertTrue("aws.sns.topic_arn" in sns_success_attributes)
+        self.assertEqual(sns_success_attributes["aws.sns.topic_arn"], _TOPIC_ARN)
 
 
 def _do_extract_kinesis_attributes() -> Dict[str, str]:
@@ -135,10 +137,30 @@ def _do_extract_sns_attributes() -> Dict[str, str]:
     return _do_extract_attributes(service_name, params)
 
 
+def _do_sns_on_success() -> Dict[str, str]:
+    service_name: str = "sns"
+    result: Dict[str, Any] = {"TopicArn": _TOPIC_ARN}
+    return _do_on_success(service_name, result)
+
+
 def _do_extract_attributes(service_name: str, params: Dict[str, str]) -> Dict[str, str]:
     mock_call_context: MagicMock = MagicMock()
     mock_call_context.params = params
     attributes: Dict[str, str] = {}
-    sqs_extension = _KNOWN_EXTENSIONS[service_name]()(mock_call_context)
-    sqs_extension.extract_attributes(attributes)
+    extension = _KNOWN_EXTENSIONS[service_name]()(mock_call_context)
+    extension.extract_attributes(attributes)
     return attributes
+
+
+def _do_on_success(service_name: str, result: Dict[str, Any]) -> Dict[str, str]:
+    span_mock: Span = MagicMock()
+    span_attributes: Dict[str, str] = {}
+
+    def set_side_effect(set_key, set_value):
+        span_attributes[set_key] = set_value
+
+    span_mock.set_attribute.side_effect = set_side_effect
+    extension = _KNOWN_EXTENSIONS[service_name]()(span_mock)
+    extension.on_success(span_mock, result)
+
+    return span_attributes

contract-tests/tests/test/amazon/botocore/botocore_test.py

@@ -467,6 +467,9 @@ def _assert_aws_span_attributes(self, resource_scope_spans: List[ResourceScopeSp
         for resource_scope_span in resource_scope_spans:
             # pylint: disable=no-member
             if resource_scope_span.span.kind in (Span.SPAN_KIND_CLIENT, Span.SPAN_KIND_PRODUCER):
+                self.assertEqual(
+                    resource_scope_span.span.kind, Span.SpanKind.Value("SPAN_KIND_" + kwargs.get("span_kind"))
+                )
                 target_spans.append(resource_scope_span.span)
         self.assertEqual(len(target_spans), 1)
         self._assert_aws_attributes(
@@ -509,6 +512,9 @@ def _assert_semantic_conventions_span_attributes(
         for resource_scope_span in resource_scope_spans:
             # pylint: disable=no-member
             if resource_scope_span.span.kind in (Span.SPAN_KIND_CLIENT, Span.SPAN_KIND_PRODUCER):
+                self.assertEqual(
+                    resource_scope_span.span.kind, Span.SpanKind.Value("SPAN_KIND_" + kwargs.get("span_kind"))
+                )
                 target_spans.append(resource_scope_span.span)
 
         self.assertEqual(len(target_spans), 1)
@@ -531,6 +537,7 @@ def _assert_semantic_conventions_attributes(
         request_specific_attributes: dict,
     ) -> None:
         attributes_dict: Dict[str, AnyValue] = self._get_attributes_dict(attributes_list)
+
         self._assert_str_attribute(attributes_dict, SpanAttributes.RPC_METHOD, operation)
         self._assert_str_attribute(attributes_dict, SpanAttributes.RPC_SYSTEM, "aws-api")
         self._assert_str_attribute(attributes_dict, SpanAttributes.RPC_SERVICE, service.split("::")[-1])