From 68c4dbcb05aa7fc312ff180130d4c1e7837a7dc3 Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Tue, 28 May 2024 16:22:03 -0700
Subject: [PATCH 01/17] Add SNS support.

---
 .../distro/_aws_attribute_keys.py             |   1 +
 .../distro/_aws_metric_attribute_generator.py |   5 +
 .../distro/patches/_botocore_patches.py       |  24 +-
 .../distro/patches/test_botocore_sns.py       |  36 +++
 .../test_aws_metric_attribute_generator.py    |   7 +
 .../distro/test_instrumentation_patch.py      |  28 ++-
 .../applications/botocore/botocore_server.py  |  49 +++-
 .../test/amazon/botocore/botocore_test.py     | 233 ++++--------------
 8 files changed, 188 insertions(+), 195 deletions(-)
 create mode 100644 aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py

diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py
index f6498ac76..5e8c8f422 100644
--- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py
+++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py
@@ -16,3 +16,4 @@
 AWS_QUEUE_URL: str = "aws.sqs.queue_url"
 AWS_QUEUE_NAME: str = "aws.sqs.queue_name"
 AWS_STREAM_NAME: str = "aws.kinesis.stream_name"
+AWS_TOPIC_ARN: str = "aws.sns.topic_arn"
diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py
index 577d28f63..5493b0263 100644
--- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py
+++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py
@@ -16,6 +16,7 @@
     AWS_REMOTE_SERVICE,
     AWS_SPAN_KIND,
     AWS_STREAM_NAME,
+    AWS_TOPIC_ARN,
 )
 from amazon.opentelemetry.distro._aws_span_processing_util import (
     LOCAL_ROOT,
@@ -78,6 +79,7 @@
 _NORMALIZED_KINESIS_SERVICE_NAME: str = "AWS::Kinesis"
 _NORMALIZED_S3_SERVICE_NAME: str = "AWS::S3"
 _NORMALIZED_SQS_SERVICE_NAME: str = "AWS::SQS"
+_NORMALIZED_SNS_SERVICE_NAME: str = "AWS::SNS"
 _DB_CONNECTION_STRING_TYPE: str = "DB::Connection"
 
 # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
@@ -372,6 +374,9 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
             remote_resource_identifier = _escape_delimiters(
                 SqsUrlParser.get_queue_name(span.attributes.get(AWS_QUEUE_URL))
             )
+        elif is_key_present(span, AWS_TOPIC_ARN):
+            remote_resource_type = _NORMALIZED_SNS_SERVICE_NAME + "::TopicArn"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_TOPIC_ARN))
     elif is_db_span(span):
         remote_resource_type = _DB_CONNECTION_STRING_TYPE
         remote_resource_identifier = _get_db_connection(span)
diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
index cf73fb345..1d11abab5 100644
--- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
+++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
@@ -4,6 +4,7 @@
 import importlib
 
 from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS
+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
 from opentelemetry.semconv.trace import SpanAttributes
@@ -12,11 +13,12 @@
 def _apply_botocore_instrumentation_patches() -> None:
     """Botocore instrumentation patches
 
-    Adds patches to provide additional support and Java parity for Kinesis, S3, and SQS.
+    Adds patches to provide additional support for Kinesis, S3, SQS and SNS.
     """
     _apply_botocore_kinesis_patch()
     _apply_botocore_s3_patch()
     _apply_botocore_sqs_patch()
+    _apply_botocore_sns_patch()
 
 
 def _apply_botocore_kinesis_patch() -> None:
@@ -65,6 +67,26 @@ def patch_extract_attributes(self, attributes: _AttributeMapT):
     _SqsExtension.extract_attributes = patch_extract_attributes
 
 
+def _apply_botocore_sns_patch() -> None:
+    """Botocore instrumentation patch for SNS
+
+    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 only be assigned with TargetArn as well.
+    """
+    old_extract_attributes = _SnsExtension.extract_attributes
+
+    def patch_extract_attributes(self, attributes: _AttributeMapT):
+        old_extract_attributes(self, attributes)
+        topic_arn = self._call_context.params.get("TopicArn")
+        if topic_arn:
+            attributes["aws.sns.topic_arn"] = topic_arn
+
+    _SnsExtension.extract_attributes = patch_extract_attributes
+
+
 # The OpenTelemetry Authors code
 def _lazy_load(module, cls):
     """Clone of upstream opentelemetry.instrumentation.botocore.extensions.lazy_load
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
new file mode 100644
index 000000000..bdaf12392
--- /dev/null
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
@@ -0,0 +1,36 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import botocore.session
+from moto import mock_aws
+
+from amazon.opentelemetry.distro.patches._instrumentation_patch import apply_instrumentation_patches
+from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
+from opentelemetry.test.test_base import TestBase
+
+
+class TestSnsExtension(TestBase):
+    def setUp(self):
+        super().setUp()
+        BotocoreInstrumentor().instrument()
+        # Apply patches
+        apply_instrumentation_patches()
+
+        session = botocore.session.get_session()
+        session.set_credentials(access_key="access-key", secret_key="secret-key")
+        self.client = session.create_client("sns", region_name="us-west-2")
+        self.topic_name = "my-topic"
+
+    def tearDown(self):
+        super().tearDown()
+        BotocoreInstrumentor().uninstrument()
+
+    @mock_aws
+    def test_create_and_delete_topic(self):
+        self.memory_exporter.clear()
+        response = self.client.create_topic(Name=self.topic_name)
+        topic_arn = response["TopicArn"]
+        self.client.delete_topic(TopicArn=topic_arn)
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(2, len(spans))
+        span = spans[1]
+        self.assertEqual(topic_arn, span.attributes["aws.sns.topic_arn"])
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py
index 072e6eeb0..49be9c9cd 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py
@@ -19,6 +19,7 @@
     AWS_REMOTE_SERVICE,
     AWS_SPAN_KIND,
     AWS_STREAM_NAME,
+    AWS_TOPIC_ARN,
 )
 from amazon.opentelemetry.distro._aws_metric_attribute_generator import _AwsMetricAttributeGenerator
 from amazon.opentelemetry.distro.metric_attribute_generator import DEPENDENCY_METRIC, SERVICE_METRIC
@@ -821,6 +822,7 @@ def test_normalize_remote_service_name_aws_sdk(self):
         self.validate_aws_sdk_service_normalization("Kinesis", "AWS::Kinesis")
         self.validate_aws_sdk_service_normalization("S3", "AWS::S3")
         self.validate_aws_sdk_service_normalization("SQS", "AWS::SQS")
+        self.validate_aws_sdk_service_normalization("SNS", "AWS::SNS")
 
     def validate_aws_sdk_service_normalization(self, service_name: str, expected_remote_service: str):
         self._mock_attribute([SpanAttributes.RPC_SYSTEM, SpanAttributes.RPC_SERVICE], ["aws-api", service_name])
@@ -977,6 +979,11 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         self._validate_remote_resource_attributes("AWS::DynamoDB::Table", "aws_table^^name")
         self._mock_attribute([SpanAttributes.AWS_DYNAMODB_TABLE_NAMES], [None])
 
+        # Validate behaviour of AWS_TOPIC_ARN attribute, then remove it
+        self._mock_attribute([AWS_TOPIC_ARN], ["arn:aws:sns:us-west-2:012345678901:test_topic"], keys, values)
+        self._validate_remote_resource_attributes("AWS::SNS::TopicArn", "arn:aws:sns:us-west-2:012345678901:test_topic")
+        self._mock_attribute([AWS_TOPIC_ARN], [None])
+
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
 
     def test_client_db_span_with_remote_resource_attributes(self):
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
index bc6e851a9..d7172874c 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
@@ -14,6 +14,7 @@
 _BUCKET_NAME: str = "bucketName"
 _QUEUE_NAME: str = "queueName"
 _QUEUE_URL: str = "queueUrl"
+_TOPIC_ARN: str = "topicArn"
 
 
 class TestInstrumentationPatch(TestCase):
@@ -69,10 +70,17 @@ def _validate_unpatched_botocore_instrumentation(self):
 
         # SQS
         self.assertTrue("sqs" in _KNOWN_EXTENSIONS, "Upstream has removed the SQS extension")
-        attributes: Dict[str, str] = _do_extract_sqs_attributes()
-        self.assertTrue("aws.queue_url" in attributes)
-        self.assertFalse("aws.sqs.queue_url" in attributes)
-        self.assertFalse("aws.sqs.queue_name" in attributes)
+        sqs_attributes: Dict[str, str] = _do_extract_sqs_attributes()
+        self.assertTrue("aws.queue_url" in sqs_attributes)
+        self.assertFalse("aws.sqs.queue_url" in sqs_attributes)
+        self.assertFalse("aws.sqs.queue_name" in sqs_attributes)
+
+        # 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):
         # Kinesis
@@ -96,6 +104,12 @@ def _validate_patched_botocore_instrumentation(self):
         self.assertTrue("aws.sqs.queue_name" in sqs_attributes)
         self.assertEqual(sqs_attributes["aws.sqs.queue_name"], _QUEUE_NAME)
 
+        # SNS
+        self.assertTrue("sns" in _KNOWN_EXTENSIONS)
+        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)
+
 
 def _do_extract_kinesis_attributes() -> Dict[str, str]:
     service_name: str = "kinesis"
@@ -115,6 +129,12 @@ def _do_extract_sqs_attributes() -> Dict[str, str]:
     return _do_extract_attributes(service_name, params)
 
 
+def _do_extract_sns_attributes() -> Dict[str, str]:
+    service_name: str = "sns"
+    params: Dict[str, str] = {"TopicArn": _TOPIC_ARN}
+    return _do_extract_attributes(service_name, params)
+
+
 def _do_extract_attributes(service_name: str, params: Dict[str, str]) -> Dict[str, str]:
     mock_call_context: MagicMock = MagicMock()
     mock_call_context.params = params
diff --git a/contract-tests/images/applications/botocore/botocore_server.py b/contract-tests/images/applications/botocore/botocore_server.py
index 192093d1f..bff95ca6b 100644
--- a/contract-tests/images/applications/botocore/botocore_server.py
+++ b/contract-tests/images/applications/botocore/botocore_server.py
@@ -41,12 +41,14 @@ def do_GET(self):
             self._handle_sqs_request()
         if self.in_path("kinesis"):
             self._handle_kinesis_request()
+        if self.in_path("sns"):
+            self._handle_sns_request()
 
         self._end_request(self.main_status)
 
     # pylint: disable=invalid-name
     def do_POST(self):
-        if self.in_path("sqserror"):
+        if self.in_path("sqserror") or self.in_path("snserror"):
             self.send_response(self.main_status)
             self.send_header("Content-type", "text/xml")
             self.end_headers()
@@ -203,6 +205,47 @@ def _handle_kinesis_request(self) -> None:
         else:
             set_main_status(404)
 
+    def _handle_sns_request(self) -> None:
+        sns_client: BaseClient = boto3.client("sns", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
+        if self.in_path(_ERROR):
+            set_main_status(400)
+            try:
+                error_client: BaseClient = boto3.client(
+                    "sns", endpoint_url=_ERROR_ENDPOINT + "/snserror", region_name=_AWS_REGION
+                )
+                topic_arn = "arn:aws:sns:us-west-2:000000000000:test_topic/snserror"
+                message = "Hello from Amazon SNS!"
+                subject = "Test Message"
+                message_attributes = {"Attribute1": {"DataType": "String", "StringValue": "Value1"}}
+                error_client.publish(
+                    TopicArn=topic_arn, Message=message, Subject=subject, MessageAttributes=message_attributes
+                )
+            except Exception as exception:
+                print("Expected exception occurred", exception)
+        elif self.in_path(_FAULT):
+            set_main_status(500)
+            try:
+                fault_client: BaseClient = boto3.client(
+                    "sns", endpoint_url=_FAULT_ENDPOINT, region_name=_AWS_REGION, config=_NO_RETRY_CONFIG
+                )
+                fault_client.get_topic_attributes(TopicArn="invalid_topic_arn")
+            except Exception as exception:
+                print("Expected exception occurred", exception)
+        elif self.in_path("gettopattributes/get-topic-attributes"):
+            set_main_status(200)
+            sns_client.get_topic_attributes(TopicArn="arn:aws:sns:us-west-2:000000000000:test_topic")
+        elif self.in_path("publishmessage/publish-message/some-message"):
+            set_main_status(200)
+            topic_arn = "arn:aws:sns:us-west-2:000000000000:test_topic"
+            message = "Hello from Amazon SNS!"
+            subject = "Test Message"
+            message_attributes = {"Attribute1": {"DataType": "String", "StringValue": "Value1"}}
+            sns_client.publish(
+                TopicArn=topic_arn, Message=message, Subject=subject, MessageAttributes=message_attributes
+            )
+        else:
+            set_main_status(404)
+
     def _end_request(self, status_code: int):
         self.send_response_only(status_code)
         self.end_headers()
@@ -247,6 +290,10 @@ def prepare_aws_server() -> None:
         # Set up Kinesis so tests can access a stream.
         kinesis_client: BaseClient = boto3.client("kinesis", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
         kinesis_client.create_stream(StreamName="test_stream", ShardCount=1)
+
+        # Set up SNS so tests can access a topic.
+        sns_client: BaseClient = boto3.client("sns", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
+        sns_client.create_topic(Name="test_topic")
     except Exception as exception:
         print("Unexpected exception occurred", exception)
 
diff --git a/contract-tests/tests/test/amazon/botocore/botocore_test.py b/contract-tests/tests/test/amazon/botocore/botocore_test.py
index f309f343b..96a08b700 100644
--- a/contract-tests/tests/test/amazon/botocore/botocore_test.py
+++ b/contract-tests/tests/test/amazon/botocore/botocore_test.py
@@ -29,6 +29,7 @@
 _AWS_QUEUE_URL: str = "aws.sqs.queue_url"
 _AWS_QUEUE_NAME: str = "aws.sqs.queue_name"
 _AWS_STREAM_NAME: str = "aws.kinesis.stream_name"
+_AWS_TOPIC_ARN: str = "aws.sns.topic_arn"
 
 
 # pylint: disable=too-many-public-methods
@@ -97,6 +98,7 @@ def test_s3_create_bucket(self):
                 SpanAttributes.AWS_S3_BUCKET: "test-bucket-name",
             },
             span_name="S3.CreateBucket",
+            span_kind="CLIENT",
         )
 
     def test_s3_create_object(self):
@@ -114,6 +116,7 @@ def test_s3_create_object(self):
                 SpanAttributes.AWS_S3_BUCKET: "test-put-object-bucket-name",
             },
             span_name="S3.PutObject",
+            span_kind="CLIENT",
         )
 
     def test_s3_get_object(self):
@@ -131,6 +134,7 @@ def test_s3_get_object(self):
                 SpanAttributes.AWS_S3_BUCKET: "test-get-object-bucket-name",
             },
             span_name="S3.GetObject",
+            span_kind="CLIENT",
         )
 
     def test_s3_error(self):
@@ -148,227 +152,79 @@ def test_s3_error(self):
                 SpanAttributes.AWS_S3_BUCKET: "-",
             },
             span_name="S3.CreateBucket",
+            span_kind="CLIENT",
         )
 
-    def test_s3_fault(self):
+    def test_sns_get_topic_attributes(self):
         self.do_test_requests(
-            "s3/fault",
-            "GET",
-            500,
-            0,
-            1,
-            remote_service="AWS::S3",
-            remote_operation="CreateBucket",
-            remote_resource_type="AWS::S3::Bucket",
-            remote_resource_identifier="valid-bucket-name",
-            request_specific_attributes={
-                SpanAttributes.AWS_S3_BUCKET: "valid-bucket-name",
-            },
-            span_name="S3.CreateBucket",
-        )
-
-    def test_dynamodb_create_table(self):
-        self.do_test_requests(
-            "ddb/createtable/some-table",
-            "GET",
-            200,
-            0,
-            0,
-            remote_service="AWS::DynamoDB",
-            remote_operation="CreateTable",
-            remote_resource_type="AWS::DynamoDB::Table",
-            remote_resource_identifier="test_table",
-            request_specific_attributes={
-                SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["test_table"],
-            },
-            span_name="DynamoDB.CreateTable",
-        )
-
-    def test_dynamodb_put_item(self):
-        self.do_test_requests(
-            "ddb/putitem/putitem-table/key",
-            "GET",
-            200,
-            0,
-            0,
-            remote_service="AWS::DynamoDB",
-            remote_operation="PutItem",
-            remote_resource_type="AWS::DynamoDB::Table",
-            remote_resource_identifier="put_test_table",
-            request_specific_attributes={
-                SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["put_test_table"],
-            },
-            span_name="DynamoDB.PutItem",
-        )
-
-    def test_dynamodb_error(self):
-        self.do_test_requests(
-            "ddb/error",
-            "GET",
-            400,
-            1,
-            0,
-            remote_service="AWS::DynamoDB",
-            remote_operation="PutItem",
-            remote_resource_type="AWS::DynamoDB::Table",
-            remote_resource_identifier="invalid_table",
-            request_specific_attributes={
-                SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["invalid_table"],
-            },
-            span_name="DynamoDB.PutItem",
-        )
-
-    def test_dynamodb_fault(self):
-        self.do_test_requests(
-            "ddb/fault",
-            "GET",
-            500,
-            0,
-            1,
-            remote_service="AWS::DynamoDB",
-            remote_operation="PutItem",
-            remote_resource_type="AWS::DynamoDB::Table",
-            remote_resource_identifier="invalid_table",
-            request_specific_attributes={
-                SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["invalid_table"],
-            },
-            span_name="DynamoDB.PutItem",
-        )
-
-    def test_sqs_create_queue(self):
-        self.do_test_requests(
-            "sqs/createqueue/some-queue",
-            "GET",
-            200,
-            0,
-            0,
-            remote_service="AWS::SQS",
-            remote_operation="CreateQueue",
-            remote_resource_type="AWS::SQS::Queue",
-            remote_resource_identifier="test_queue",
-            request_specific_attributes={
-                _AWS_QUEUE_NAME: "test_queue",
-            },
-            span_name="SQS.CreateQueue",
-        )
-
-    def test_sqs_send_message(self):
-        self.do_test_requests(
-            "sqs/publishqueue/some-queue",
+            "sns/gettopattributes/get-topic-attributes",
             "GET",
             200,
             0,
             0,
-            remote_service="AWS::SQS",
-            remote_operation="SendMessage",
-            remote_resource_type="AWS::SQS::Queue",
-            remote_resource_identifier="test_put_get_queue",
+            remote_service="AWS::SNS",
+            remote_operation="GetTopicAttributes",
+            remote_resource_type="AWS::SNS::TopicArn",
+            remote_resource_identifier="arn:aws:sns:us-west-2:000000000000:test_topic",
             request_specific_attributes={
-                _AWS_QUEUE_URL: "http://localstack:4566/000000000000/test_put_get_queue",
+                _AWS_TOPIC_ARN: "arn:aws:sns:us-west-2:000000000000:test_topic",
             },
-            span_name="SQS.SendMessage",
+            span_name="SNS.GetTopicAttributes",
+            span_kind="CLIENT",
         )
 
-    def test_sqs_receive_message(self):
+    def test_sns_publish_message(self):
         self.do_test_requests(
-            "sqs/consumequeue/some-queue",
+            "sns/publishmessage/publish-message/some-message",
             "GET",
             200,
             0,
             0,
-            remote_service="AWS::SQS",
-            remote_operation="ReceiveMessage",
-            remote_resource_type="AWS::SQS::Queue",
-            remote_resource_identifier="test_put_get_queue",
+            remote_service="AWS::SNS",
+            remote_operation="Publish",
+            remote_resource_type="AWS::SNS::TopicArn",
+            remote_resource_identifier="arn:aws:sns:us-west-2:000000000000:test_topic",
             request_specific_attributes={
-                _AWS_QUEUE_URL: "http://localstack:4566/000000000000/test_put_get_queue",
+                _AWS_TOPIC_ARN: "arn:aws:sns:us-west-2:000000000000:test_topic",
             },
-            span_name="SQS.ReceiveMessage",
+            span_name="test_topic send",
+            span_kind="PRODUCER",
         )
 
-    def test_sqs_error(self):
+    def test_sns_error(self):
         self.do_test_requests(
-            "sqs/error",
+            "sns/error",
             "GET",
             400,
             1,
             0,
-            remote_service="AWS::SQS",
-            remote_operation="SendMessage",
-            remote_resource_type="AWS::SQS::Queue",
-            remote_resource_identifier="sqserror",
+            remote_service="AWS::SNS",
+            remote_operation="Publish",
+            remote_resource_type="AWS::SNS::TopicArn",
+            remote_resource_identifier="arn:aws:sns:us-west-2:000000000000:test_topic/snserror",
             request_specific_attributes={
-                _AWS_QUEUE_URL: "http://error.test:8080/000000000000/sqserror",
+                _AWS_TOPIC_ARN: "arn:aws:sns:us-west-2:000000000000:test_topic/snserror",
             },
-            span_name="SQS.SendMessage",
+            span_name="test_topic/snserror send",
+            span_kind="PRODUCER",
         )
 
-    def test_sqs_fault(self):
+    def test_sns_fault(self):
         self.do_test_requests(
-            "sqs/fault",
+            "sns/fault",
             "GET",
             500,
             0,
             1,
-            remote_service="AWS::SQS",
-            remote_operation="CreateQueue",
-            remote_resource_type="AWS::SQS::Queue",
-            remote_resource_identifier="invalid_test",
+            remote_service="AWS::SNS",
+            remote_operation="GetTopicAttributes",
+            remote_resource_type="AWS::SNS::TopicArn",
+            remote_resource_identifier="invalid_topic_arn",
             request_specific_attributes={
-                _AWS_QUEUE_NAME: "invalid_test",
+                _AWS_TOPIC_ARN: "invalid_topic_arn",
             },
-            span_name="SQS.CreateQueue",
-        )
-
-    def test_kinesis_put_record(self):
-        self.do_test_requests(
-            "kinesis/putrecord/my-stream",
-            "GET",
-            200,
-            0,
-            0,
-            remote_service="AWS::Kinesis",
-            remote_operation="PutRecord",
-            remote_resource_type="AWS::Kinesis::Stream",
-            remote_resource_identifier="test_stream",
-            request_specific_attributes={
-                _AWS_STREAM_NAME: "test_stream",
-            },
-            span_name="Kinesis.PutRecord",
-        )
-
-    def test_kinesis_error(self):
-        self.do_test_requests(
-            "kinesis/error",
-            "GET",
-            400,
-            1,
-            0,
-            remote_service="AWS::Kinesis",
-            remote_operation="PutRecord",
-            remote_resource_type="AWS::Kinesis::Stream",
-            remote_resource_identifier="invalid_stream",
-            request_specific_attributes={
-                _AWS_STREAM_NAME: "invalid_stream",
-            },
-            span_name="Kinesis.PutRecord",
-        )
-
-    def test_kinesis_fault(self):
-        self.do_test_requests(
-            "kinesis/fault",
-            "GET",
-            500,
-            0,
-            1,
-            remote_service="AWS::Kinesis",
-            remote_operation="PutRecord",
-            remote_resource_type="AWS::Kinesis::Stream",
-            remote_resource_identifier="test_stream",
-            request_specific_attributes={
-                _AWS_STREAM_NAME: "test_stream",
-            },
-            span_name="Kinesis.PutRecord",
+            span_name="SNS.GetTopicAttributes",
+            span_kind="CLIENT",
         )
 
     @override
@@ -376,9 +232,8 @@ def _assert_aws_span_attributes(self, resource_scope_spans: List[ResourceScopeSp
         target_spans: List[Span] = []
         for resource_scope_span in resource_scope_spans:
             # pylint: disable=no-member
-            if resource_scope_span.span.kind == Span.SPAN_KIND_CLIENT:
+            if resource_scope_span.span.kind in (Span.SPAN_KIND_CLIENT, Span.SPAN_KIND_PRODUCER):
                 target_spans.append(resource_scope_span.span)
-
         self.assertEqual(len(target_spans), 1)
         self._assert_aws_attributes(
             target_spans[0].attributes,
@@ -419,7 +274,7 @@ def _assert_semantic_conventions_span_attributes(
         target_spans: List[Span] = []
         for resource_scope_span in resource_scope_spans:
             # pylint: disable=no-member
-            if resource_scope_span.span.kind == Span.SPAN_KIND_CLIENT:
+            if resource_scope_span.span.kind in (Span.SPAN_KIND_CLIENT, Span.SPAN_KIND_PRODUCER):
                 target_spans.append(resource_scope_span.span)
 
         self.assertEqual(len(target_spans), 1)
@@ -485,7 +340,7 @@ def _assert_metric_attributes(
         self._assert_str_attribute(attribute_dict, AWS_LOCAL_OPERATION, "InternalOperation")
         self._assert_str_attribute(attribute_dict, AWS_REMOTE_SERVICE, kwargs.get("remote_service"))
         self._assert_str_attribute(attribute_dict, AWS_REMOTE_OPERATION, kwargs.get("remote_operation"))
-        self._assert_str_attribute(attribute_dict, AWS_SPAN_KIND, "CLIENT")
+        self._assert_str_attribute(attribute_dict, AWS_SPAN_KIND, kwargs.get("span_kind"))
         remote_resource_type = kwargs.get("remote_resource_type", "None")
         remote_resource_identifier = kwargs.get("remote_resource_identifier", "None")
         if remote_resource_type != "None":

From 1df99ebdcd937bbd03d081bcb434e5d258c08acc Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 09:56:51 -0700
Subject: [PATCH 02/17] Add back contract test.

---
 .../test/amazon/botocore/botocore_test.py     | 234 ++++++++++++++++++
 tox.ini                                       |   2 +-
 2 files changed, 235 insertions(+), 1 deletion(-)

diff --git a/contract-tests/tests/test/amazon/botocore/botocore_test.py b/contract-tests/tests/test/amazon/botocore/botocore_test.py
index 96a08b700..2894e8ba4 100644
--- a/contract-tests/tests/test/amazon/botocore/botocore_test.py
+++ b/contract-tests/tests/test/amazon/botocore/botocore_test.py
@@ -155,6 +155,240 @@ def test_s3_error(self):
             span_kind="CLIENT",
         )
 
+    def test_s3_fault(self):
+        self.do_test_requests(
+            "s3/fault",
+            "GET",
+            500,
+            0,
+            1,
+            remote_service="AWS::S3",
+            remote_operation="CreateBucket",
+            remote_resource_type="AWS::S3::Bucket",
+            remote_resource_identifier="valid-bucket-name",
+            request_specific_attributes={
+                SpanAttributes.AWS_S3_BUCKET: "valid-bucket-name",
+            },
+            span_name="S3.CreateBucket",
+            span_kind="CLIENT",
+        )
+
+    def test_dynamodb_create_table(self):
+        self.do_test_requests(
+            "ddb/createtable/some-table",
+            "GET",
+            200,
+            0,
+            0,
+            remote_service="AWS::DynamoDB",
+            remote_operation="CreateTable",
+            remote_resource_type="AWS::DynamoDB::Table",
+            remote_resource_identifier="test_table",
+            request_specific_attributes={
+                SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["test_table"],
+            },
+            span_name="DynamoDB.CreateTable",
+            span_kind="CLIENT",
+        )
+
+    def test_dynamodb_put_item(self):
+        self.do_test_requests(
+            "ddb/putitem/putitem-table/key",
+            "GET",
+            200,
+            0,
+            0,
+            remote_service="AWS::DynamoDB",
+            remote_operation="PutItem",
+            remote_resource_type="AWS::DynamoDB::Table",
+            remote_resource_identifier="put_test_table",
+            request_specific_attributes={
+                SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["put_test_table"],
+            },
+            span_name="DynamoDB.PutItem",
+            span_kind="CLIENT",
+        )
+
+    def test_dynamodb_error(self):
+        self.do_test_requests(
+            "ddb/error",
+            "GET",
+            400,
+            1,
+            0,
+            remote_service="AWS::DynamoDB",
+            remote_operation="PutItem",
+            remote_resource_type="AWS::DynamoDB::Table",
+            remote_resource_identifier="invalid_table",
+            request_specific_attributes={
+                SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["invalid_table"],
+            },
+            span_name="DynamoDB.PutItem",
+            span_kind="CLIENT",
+        )
+
+    def test_dynamodb_fault(self):
+        self.do_test_requests(
+            "ddb/fault",
+            "GET",
+            500,
+            0,
+            1,
+            remote_service="AWS::DynamoDB",
+            remote_operation="PutItem",
+            remote_resource_type="AWS::DynamoDB::Table",
+            remote_resource_identifier="invalid_table",
+            request_specific_attributes={
+                SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["invalid_table"],
+            },
+            span_name="DynamoDB.PutItem",
+            span_kind="CLIENT",
+        )
+
+    def test_sqs_create_queue(self):
+        self.do_test_requests(
+            "sqs/createqueue/some-queue",
+            "GET",
+            200,
+            0,
+            0,
+            remote_service="AWS::SQS",
+            remote_operation="CreateQueue",
+            remote_resource_type="AWS::SQS::Queue",
+            remote_resource_identifier="test_queue",
+            request_specific_attributes={
+                _AWS_QUEUE_NAME: "test_queue",
+            },
+            span_name="SQS.CreateQueue",
+            span_kind="CLIENT",
+        )
+
+    def test_sqs_send_message(self):
+        self.do_test_requests(
+            "sqs/publishqueue/some-queue",
+            "GET",
+            200,
+            0,
+            0,
+            remote_service="AWS::SQS",
+            remote_operation="SendMessage",
+            remote_resource_type="AWS::SQS::Queue",
+            remote_resource_identifier="test_put_get_queue",
+            request_specific_attributes={
+                _AWS_QUEUE_URL: "http://localstack:4566/000000000000/test_put_get_queue",
+            },
+            span_name="SQS.SendMessage",
+            span_kind="CLIENT",
+        )
+
+    def test_sqs_receive_message(self):
+        self.do_test_requests(
+            "sqs/consumequeue/some-queue",
+            "GET",
+            200,
+            0,
+            0,
+            remote_service="AWS::SQS",
+            remote_operation="ReceiveMessage",
+            remote_resource_type="AWS::SQS::Queue",
+            remote_resource_identifier="test_put_get_queue",
+            request_specific_attributes={
+                _AWS_QUEUE_URL: "http://localstack:4566/000000000000/test_put_get_queue",
+            },
+            span_name="SQS.ReceiveMessage",
+            span_kind="CLIENT",
+        )
+
+    def test_sqs_error(self):
+        self.do_test_requests(
+            "sqs/error",
+            "GET",
+            400,
+            1,
+            0,
+            remote_service="AWS::SQS",
+            remote_operation="SendMessage",
+            remote_resource_type="AWS::SQS::Queue",
+            remote_resource_identifier="sqserror",
+            request_specific_attributes={
+                _AWS_QUEUE_URL: "http://error.test:8080/000000000000/sqserror",
+            },
+            span_name="SQS.SendMessage",
+            span_kind="CLIENT",
+        )
+
+    def test_sqs_fault(self):
+        self.do_test_requests(
+            "sqs/fault",
+            "GET",
+            500,
+            0,
+            1,
+            remote_service="AWS::SQS",
+            remote_operation="CreateQueue",
+            remote_resource_type="AWS::SQS::Queue",
+            remote_resource_identifier="invalid_test",
+            request_specific_attributes={
+                _AWS_QUEUE_NAME: "invalid_test",
+            },
+            span_name="SQS.CreateQueue",
+            span_kind="CLIENT",
+        )
+
+    def test_kinesis_put_record(self):
+        self.do_test_requests(
+            "kinesis/putrecord/my-stream",
+            "GET",
+            200,
+            0,
+            0,
+            remote_service="AWS::Kinesis",
+            remote_operation="PutRecord",
+            remote_resource_type="AWS::Kinesis::Stream",
+            remote_resource_identifier="test_stream",
+            request_specific_attributes={
+                _AWS_STREAM_NAME: "test_stream",
+            },
+            span_name="Kinesis.PutRecord",
+            span_kind="CLIENT",
+        )
+
+    def test_kinesis_error(self):
+        self.do_test_requests(
+            "kinesis/error",
+            "GET",
+            400,
+            1,
+            0,
+            remote_service="AWS::Kinesis",
+            remote_operation="PutRecord",
+            remote_resource_type="AWS::Kinesis::Stream",
+            remote_resource_identifier="invalid_stream",
+            request_specific_attributes={
+                _AWS_STREAM_NAME: "invalid_stream",
+            },
+            span_name="Kinesis.PutRecord",
+            span_kind="CLIENT",
+        )
+
+    def test_kinesis_fault(self):
+        self.do_test_requests(
+            "kinesis/fault",
+            "GET",
+            500,
+            0,
+            1,
+            remote_service="AWS::Kinesis",
+            remote_operation="PutRecord",
+            remote_resource_type="AWS::Kinesis::Stream",
+            remote_resource_identifier="test_stream",
+            request_specific_attributes={
+                _AWS_STREAM_NAME: "test_stream",
+            },
+            span_name="Kinesis.PutRecord",
+            span_kind="CLIENT",
+        )
+
     def test_sns_get_topic_attributes(self):
         self.do_test_requests(
             "sns/gettopattributes/get-topic-attributes",
diff --git a/tox.ini b/tox.ini
index 7cd36fbbf..d167e635d 100644
--- a/tox.ini
+++ b/tox.ini
@@ -26,7 +26,7 @@ changedir =
 
 commands_pre =
 ; Install without -e to test the actual installation
-  3.{7,8,9,10,11}: python -m pip install -U pip setuptools wheel
+  3.{7,8,9,10,11}: python -m pip install -U pip setuptools wheel moto
 ; Install common packages for all the tests. These are not needed in all the
 ; cases but it saves a lot of boilerplate in this file.
   test: pip install botocore

From c4838d04c450ac332c9241f70410d3b876993bac Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 10:02:56 -0700
Subject: [PATCH 03/17] Add test dependencies.

---
 tox.ini | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tox.ini b/tox.ini
index d167e635d..640e880f5 100644
--- a/tox.ini
+++ b/tox.ini
@@ -34,6 +34,7 @@ commands_pre =
   test: pip install "opentelemetry-sdk[test] @ {env:CORE_REPO}#egg=opentelemetry-sdk&subdirectory=opentelemetry-sdk"
   test: pip install "opentelemetry-instrumentation[test] @ {env:CONTRIB_REPO}#egg=opentelemetry-instrumentation&subdirectory=opentelemetry-instrumentation"
   test: pip install "opentelemetry-exporter-otlp[test] @ {env:CORE_REPO}#egg=opentelemetry-exporter-otlp&subdirectory=exporter/opentelemetry-exporter-otlp"
+  test: pip install "opentelemetry-test-utils[test] @ {env:CORE_REPO}#egg=opentelemetry-test-utils&subdirectory=tests/opentelemetry-test-utils"
   aws-opentelemetry-distro: pip install {toxinidir}/aws-opentelemetry-distro
 
 commands =

From a2738c8d61be49590b42dc9abd85751204efb61f Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 10:06:06 -0700
Subject: [PATCH 04/17] disable pylint check.

---
 .../amazon/opentelemetry/distro/patches/test_botocore_sns.py    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
index bdaf12392..dbb1cb855 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
@@ -7,7 +7,7 @@
 from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
 from opentelemetry.test.test_base import TestBase
 
-
+# pylint: disable=C0103
 class TestSnsExtension(TestBase):
     def setUp(self):
         super().setUp()

From f5804d9cbc85f70d7cfee9abdf96018a3b43f5d6 Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 10:15:39 -0700
Subject: [PATCH 05/17] Apply lint

---
 .../amazon/opentelemetry/distro/patches/test_botocore_sns.py     | 1 +
 1 file changed, 1 insertion(+)

diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
index dbb1cb855..b136c84e6 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
@@ -7,6 +7,7 @@
 from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
 from opentelemetry.test.test_base import TestBase
 
+
 # pylint: disable=C0103
 class TestSnsExtension(TestBase):
     def setUp(self):

From f4bf8a69631d04750f43f8d2801bf3e645384b6e Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 13:30:18 -0700
Subject: [PATCH 06/17] Add reponse instrument.

---
 .../distro/patches/_botocore_patches.py         | 17 +++++++++++++++--
 .../distro/patches/test_botocore_sns.py         |  5 +++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
index 1d11abab5..efc5af490 100644
--- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
+++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
@@ -6,8 +6,9 @@
 from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS
 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
+from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension, _BotoResultT
 from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace.span import Span
 
 
 def _apply_botocore_instrumentation_patches() -> None:
@@ -74,7 +75,7 @@ def _apply_botocore_sns_patch() -> None:
     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 only be assigned with TargetArn as well.
+    but we are not using it as it can be assigned with TargetArn as well.
     """
     old_extract_attributes = _SnsExtension.extract_attributes
 
@@ -84,7 +85,19 @@ def patch_extract_attributes(self, attributes: _AttributeMapT):
         if topic_arn:
             attributes["aws.sns.topic_arn"] = topic_arn
 
+    old_on_success = _SnsExtension.on_success
+
+    def patch_on_success(self, span: Span, result: _BotoResultT):
+        old_on_success(self, span, result)
+        topic_arn = result.get("TopicArn")
+        if topic_arn:
+            span.set_attribute(
+                "aws.sns.topic_arn",
+                topic_arn,
+            )
+
     _SnsExtension.extract_attributes = patch_extract_attributes
+    _SnsExtension.on_success = patch_on_success
 
 
 # The OpenTelemetry Authors code
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
index b136c84e6..d84c8f92f 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
@@ -30,6 +30,11 @@ def test_create_and_delete_topic(self):
         self.memory_exporter.clear()
         response = self.client.create_topic(Name=self.topic_name)
         topic_arn = response["TopicArn"]
+        create_spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(1, len(create_spans))
+        create_span = create_spans[0]
+        self.assertEqual(topic_arn, create_span.attributes["aws.sns.topic_arn"])
+
         self.client.delete_topic(TopicArn=topic_arn)
         spans = self.memory_exporter.get_finished_spans()
         self.assertEqual(2, len(spans))

From 2f7f1e966ae95383023dfb0f267df4a4a6ff6780 Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 14:56:45 -0700
Subject: [PATCH 07/17] Address comments.

---
 .../distro/patches/test_botocore_sns.py       | 42 -------------------
 .../distro/test_instrumentation_patch.py      |  2 -
 .../test/amazon/botocore/botocore_test.py     |  3 ++
 3 files changed, 3 insertions(+), 44 deletions(-)
 delete mode 100644 aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py

diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
deleted file mode 100644
index d84c8f92f..000000000
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_botocore_sns.py
+++ /dev/null
@@ -1,42 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-import botocore.session
-from moto import mock_aws
-
-from amazon.opentelemetry.distro.patches._instrumentation_patch import apply_instrumentation_patches
-from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
-from opentelemetry.test.test_base import TestBase
-
-
-# pylint: disable=C0103
-class TestSnsExtension(TestBase):
-    def setUp(self):
-        super().setUp()
-        BotocoreInstrumentor().instrument()
-        # Apply patches
-        apply_instrumentation_patches()
-
-        session = botocore.session.get_session()
-        session.set_credentials(access_key="access-key", secret_key="secret-key")
-        self.client = session.create_client("sns", region_name="us-west-2")
-        self.topic_name = "my-topic"
-
-    def tearDown(self):
-        super().tearDown()
-        BotocoreInstrumentor().uninstrument()
-
-    @mock_aws
-    def test_create_and_delete_topic(self):
-        self.memory_exporter.clear()
-        response = self.client.create_topic(Name=self.topic_name)
-        topic_arn = response["TopicArn"]
-        create_spans = self.memory_exporter.get_finished_spans()
-        self.assertEqual(1, len(create_spans))
-        create_span = create_spans[0]
-        self.assertEqual(topic_arn, create_span.attributes["aws.sns.topic_arn"])
-
-        self.client.delete_topic(TopicArn=topic_arn)
-        spans = self.memory_exporter.get_finished_spans()
-        self.assertEqual(2, len(spans))
-        span = spans[1]
-        self.assertEqual(topic_arn, span.attributes["aws.sns.topic_arn"])
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
index d7172874c..454d32d65 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
@@ -78,8 +78,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):
diff --git a/contract-tests/tests/test/amazon/botocore/botocore_test.py b/contract-tests/tests/test/amazon/botocore/botocore_test.py
index 2894e8ba4..1eb845684 100644
--- a/contract-tests/tests/test/amazon/botocore/botocore_test.py
+++ b/contract-tests/tests/test/amazon/botocore/botocore_test.py
@@ -467,6 +467,7 @@ 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, kwargs.get("span_kind"))
                 target_spans.append(resource_scope_span.span)
         self.assertEqual(len(target_spans), 1)
         self._assert_aws_attributes(
@@ -509,6 +510,7 @@ 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, kwargs.get("span_kind"))
                 target_spans.append(resource_scope_span.span)
 
         self.assertEqual(len(target_spans), 1)
@@ -531,6 +533,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])

From 3e49c9f7f8c36c6326a9a9c1ca266d16e255f095 Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 16:14:21 -0700
Subject: [PATCH 08/17] Add on_success unit tests.

---
 .../distro/patches/_botocore_patches.py       |  2 --
 .../distro/test_instrumentation_patch.py      | 30 +++++++++++++++++--
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
index efc5af490..717d5f58e 100644
--- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
+++ b/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
 
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
index 454d32d65..3c4d9dd5f 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
+++ b/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"
@@ -107,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]:
@@ -133,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

From d0f9086a6dd2a81486716c4d6015563c837533f2 Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 16:31:43 -0700
Subject: [PATCH 09/17] Refine SpanKind.

---
 contract-tests/tests/test/amazon/botocore/botocore_test.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/contract-tests/tests/test/amazon/botocore/botocore_test.py b/contract-tests/tests/test/amazon/botocore/botocore_test.py
index 1eb845684..ce4db5cb7 100644
--- a/contract-tests/tests/test/amazon/botocore/botocore_test.py
+++ b/contract-tests/tests/test/amazon/botocore/botocore_test.py
@@ -22,6 +22,7 @@
 from opentelemetry.proto.metrics.v1.metrics_pb2 import ExponentialHistogramDataPoint, Metric
 from opentelemetry.proto.trace.v1.trace_pb2 import Span
 from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace import SpanKind
 
 _logger: Logger = getLogger(__name__)
 _logger.setLevel(INFO)
@@ -467,7 +468,7 @@ 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, kwargs.get("span_kind"))
+                self.assertEqual(resource_scope_span.span.kind, SpanKind[kwargs.get("span_kind")].value)
                 target_spans.append(resource_scope_span.span)
         self.assertEqual(len(target_spans), 1)
         self._assert_aws_attributes(
@@ -510,7 +511,7 @@ 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, kwargs.get("span_kind"))
+                self.assertEqual(resource_scope_span.span.kind, SpanKind[kwargs.get("span_kind")].value)
                 target_spans.append(resource_scope_span.span)
 
         self.assertEqual(len(target_spans), 1)

From 86ca4e748d3785d8c4cdff5360eb01a487cb0760 Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 16:44:49 -0700
Subject: [PATCH 10/17] Fix.

---
 contract-tests/tests/test/amazon/botocore/botocore_test.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/contract-tests/tests/test/amazon/botocore/botocore_test.py b/contract-tests/tests/test/amazon/botocore/botocore_test.py
index ce4db5cb7..4f0a9e01e 100644
--- a/contract-tests/tests/test/amazon/botocore/botocore_test.py
+++ b/contract-tests/tests/test/amazon/botocore/botocore_test.py
@@ -20,9 +20,8 @@
 )
 from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
 from opentelemetry.proto.metrics.v1.metrics_pb2 import ExponentialHistogramDataPoint, Metric
-from opentelemetry.proto.trace.v1.trace_pb2 import Span
+from opentelemetry.proto.trace.v1.trace_pb2 import Span, SpanKind
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace import SpanKind
 
 _logger: Logger = getLogger(__name__)
 _logger.setLevel(INFO)
@@ -468,7 +467,7 @@ 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, SpanKind[kwargs.get("span_kind")].value)
+                self.assertEqual(resource_scope_span.span.kind, SpanKind['SPAN_KIND_' + kwargs.get("span_kind")].value)
                 target_spans.append(resource_scope_span.span)
         self.assertEqual(len(target_spans), 1)
         self._assert_aws_attributes(
@@ -511,7 +510,7 @@ 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, SpanKind[kwargs.get("span_kind")].value)
+                self.assertEqual(resource_scope_span.span.kind, SpanKind['SPAN_KIND_' + kwargs.get("span_kind")].value)
                 target_spans.append(resource_scope_span.span)
 
         self.assertEqual(len(target_spans), 1)

From a0b05f55706e3a2461bd3514bbebc43cab7ebadc Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 16:57:39 -0700
Subject: [PATCH 11/17] Fix.

---
 contract-tests/tests/test/amazon/botocore/botocore_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/contract-tests/tests/test/amazon/botocore/botocore_test.py b/contract-tests/tests/test/amazon/botocore/botocore_test.py
index 4f0a9e01e..7f0f045aa 100644
--- a/contract-tests/tests/test/amazon/botocore/botocore_test.py
+++ b/contract-tests/tests/test/amazon/botocore/botocore_test.py
@@ -467,7 +467,7 @@ 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, SpanKind['SPAN_KIND_' + kwargs.get("span_kind")].value)
+                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(
@@ -510,7 +510,7 @@ 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, SpanKind['SPAN_KIND_' + kwargs.get("span_kind")].value)
+                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)

From e07ea8d6cb7c21b60a5f0a2a448c559280072821 Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Wed, 29 May 2024 17:05:02 -0700
Subject: [PATCH 12/17] Fix.

---
 .../tests/test/amazon/botocore/botocore_test.py        | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/contract-tests/tests/test/amazon/botocore/botocore_test.py b/contract-tests/tests/test/amazon/botocore/botocore_test.py
index 7f0f045aa..502b2a2f2 100644
--- a/contract-tests/tests/test/amazon/botocore/botocore_test.py
+++ b/contract-tests/tests/test/amazon/botocore/botocore_test.py
@@ -20,7 +20,7 @@
 )
 from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
 from opentelemetry.proto.metrics.v1.metrics_pb2 import ExponentialHistogramDataPoint, Metric
-from opentelemetry.proto.trace.v1.trace_pb2 import Span, SpanKind
+from opentelemetry.proto.trace.v1.trace_pb2 import Span
 from opentelemetry.semconv.trace import SpanAttributes
 
 _logger: Logger = getLogger(__name__)
@@ -467,7 +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")))
+                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(
@@ -510,7 +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")))
+                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)

From b0fec2ea346085f2687915cd9965e16677c33d34 Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Thu, 30 May 2024 09:12:07 -0700
Subject: [PATCH 13/17] Use Topic instead of TopicArn.

---
 .../distro/_aws_metric_attribute_generator.py             | 2 +-
 .../distro/test_aws_metric_attribute_generator.py         | 2 +-
 .../opentelemetry/distro/test_instrumentation_patch.py    | 2 ++
 .../tests/test/amazon/botocore/botocore_test.py           | 8 ++++----
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py
index 5493b0263..903746055 100644
--- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py
+++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py
@@ -375,7 +375,7 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
                 SqsUrlParser.get_queue_name(span.attributes.get(AWS_QUEUE_URL))
             )
         elif is_key_present(span, AWS_TOPIC_ARN):
-            remote_resource_type = _NORMALIZED_SNS_SERVICE_NAME + "::TopicArn"
+            remote_resource_type = _NORMALIZED_SNS_SERVICE_NAME + "::Topic"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_TOPIC_ARN))
     elif is_db_span(span):
         remote_resource_type = _DB_CONNECTION_STRING_TYPE
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py
index 49be9c9cd..400e0051d 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py
@@ -981,7 +981,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
 
         # Validate behaviour of AWS_TOPIC_ARN attribute, then remove it
         self._mock_attribute([AWS_TOPIC_ARN], ["arn:aws:sns:us-west-2:012345678901:test_topic"], keys, values)
-        self._validate_remote_resource_attributes("AWS::SNS::TopicArn", "arn:aws:sns:us-west-2:012345678901:test_topic")
+        self._validate_remote_resource_attributes("AWS::SNS::Topic", "arn:aws:sns:us-west-2:012345678901:test_topic")
         self._mock_attribute([AWS_TOPIC_ARN], [None])
 
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
index 3c4d9dd5f..dac288063 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
@@ -80,6 +80,8 @@ def _validate_unpatched_botocore_instrumentation(self):
         self.assertTrue("sns" in _KNOWN_EXTENSIONS, "Upstream has removed the SNS extension")
         sns_attributes: Dict[str, str] = _do_extract_sns_attributes()
         self.assertFalse("aws.sns.topic_arn" in sns_attributes)
+        sns_success_attributes: Dict[str, str] = _do_sns_on_success()
+        self.assertFlase("aws.sns.topic_arn" in sns_success_attributes)
 
     def _validate_patched_botocore_instrumentation(self):
         # Kinesis
diff --git a/contract-tests/tests/test/amazon/botocore/botocore_test.py b/contract-tests/tests/test/amazon/botocore/botocore_test.py
index 502b2a2f2..c1522d74d 100644
--- a/contract-tests/tests/test/amazon/botocore/botocore_test.py
+++ b/contract-tests/tests/test/amazon/botocore/botocore_test.py
@@ -398,7 +398,7 @@ def test_sns_get_topic_attributes(self):
             0,
             remote_service="AWS::SNS",
             remote_operation="GetTopicAttributes",
-            remote_resource_type="AWS::SNS::TopicArn",
+            remote_resource_type="AWS::SNS::Topic",
             remote_resource_identifier="arn:aws:sns:us-west-2:000000000000:test_topic",
             request_specific_attributes={
                 _AWS_TOPIC_ARN: "arn:aws:sns:us-west-2:000000000000:test_topic",
@@ -416,7 +416,7 @@ def test_sns_publish_message(self):
             0,
             remote_service="AWS::SNS",
             remote_operation="Publish",
-            remote_resource_type="AWS::SNS::TopicArn",
+            remote_resource_type="AWS::SNS::Topic",
             remote_resource_identifier="arn:aws:sns:us-west-2:000000000000:test_topic",
             request_specific_attributes={
                 _AWS_TOPIC_ARN: "arn:aws:sns:us-west-2:000000000000:test_topic",
@@ -434,7 +434,7 @@ def test_sns_error(self):
             0,
             remote_service="AWS::SNS",
             remote_operation="Publish",
-            remote_resource_type="AWS::SNS::TopicArn",
+            remote_resource_type="AWS::SNS::Topic",
             remote_resource_identifier="arn:aws:sns:us-west-2:000000000000:test_topic/snserror",
             request_specific_attributes={
                 _AWS_TOPIC_ARN: "arn:aws:sns:us-west-2:000000000000:test_topic/snserror",
@@ -452,7 +452,7 @@ def test_sns_fault(self):
             1,
             remote_service="AWS::SNS",
             remote_operation="GetTopicAttributes",
-            remote_resource_type="AWS::SNS::TopicArn",
+            remote_resource_type="AWS::SNS::Topic",
             remote_resource_identifier="invalid_topic_arn",
             request_specific_attributes={
                 _AWS_TOPIC_ARN: "invalid_topic_arn",

From 79c7ab21edac94ae26f503c2e6d4e4cba519dc11 Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Thu, 30 May 2024 09:14:46 -0700
Subject: [PATCH 14/17] Fix typo.

---
 .../amazon/opentelemetry/distro/test_instrumentation_patch.py   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
index dac288063..a55f2edbe 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
@@ -81,7 +81,7 @@ def _validate_unpatched_botocore_instrumentation(self):
         sns_attributes: Dict[str, str] = _do_extract_sns_attributes()
         self.assertFalse("aws.sns.topic_arn" in sns_attributes)
         sns_success_attributes: Dict[str, str] = _do_sns_on_success()
-        self.assertFlase("aws.sns.topic_arn" in sns_success_attributes)
+        self.assertFalse("aws.sns.topic_arn" in sns_success_attributes)
 
     def _validate_patched_botocore_instrumentation(self):
         # Kinesis

From d9c3cbb147ade7e6bb5d9a1ee7253ea8f41b87da Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Tue, 25 Jun 2024 11:56:52 -0700
Subject: [PATCH 15/17] Use otel upstream logic to instrument SNS.

---
 .../distro/_aws_attribute_keys.py             |  1 -
 .../distro/_aws_metric_attribute_generator.py |  9 ++-
 .../distro/patches/_botocore_patches.py       | 37 +---------
 .../test_aws_metric_attribute_generator.py    | 10 ++-
 .../distro/test_instrumentation_patch.py      | 47 +-----------
 .../applications/botocore/botocore_server.py  | 47 +-----------
 .../test/amazon/botocore/botocore_test.py     | 73 -------------------
 7 files changed, 17 insertions(+), 207 deletions(-)

diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py
index 5e8c8f422..f6498ac76 100644
--- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py
+++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py
@@ -16,4 +16,3 @@
 AWS_QUEUE_URL: str = "aws.sqs.queue_url"
 AWS_QUEUE_NAME: str = "aws.sqs.queue_name"
 AWS_STREAM_NAME: str = "aws.kinesis.stream_name"
-AWS_TOPIC_ARN: str = "aws.sns.topic_arn"
diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py
index 903746055..f3903155f 100644
--- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py
+++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py
@@ -16,7 +16,6 @@
     AWS_REMOTE_SERVICE,
     AWS_SPAN_KIND,
     AWS_STREAM_NAME,
-    AWS_TOPIC_ARN,
 )
 from amazon.opentelemetry.distro._aws_span_processing_util import (
     LOCAL_ROOT,
@@ -59,6 +58,7 @@
 _HTTP_METHOD: str = SpanAttributes.HTTP_METHOD
 _HTTP_URL: str = SpanAttributes.HTTP_URL
 _MESSAGING_OPERATION: str = SpanAttributes.MESSAGING_OPERATION
+_MESSAGING_DESTINATION: str = SpanAttributes.MESSAGING_DESTINATION
 _MESSAGING_SYSTEM: str = SpanAttributes.MESSAGING_SYSTEM
 _NET_PEER_NAME: str = SpanAttributes.NET_PEER_NAME
 _NET_PEER_PORT: str = SpanAttributes.NET_PEER_PORT
@@ -374,9 +374,12 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
             remote_resource_identifier = _escape_delimiters(
                 SqsUrlParser.get_queue_name(span.attributes.get(AWS_QUEUE_URL))
             )
-        elif is_key_present(span, AWS_TOPIC_ARN):
+        elif (
+            is_key_present(span, _MESSAGING_DESTINATION)
+            and _normalize_remote_service_name(span, _get_remote_service(span, _RPC_SERVICE)) == "AWS::SNS"
+        ):
             remote_resource_type = _NORMALIZED_SNS_SERVICE_NAME + "::Topic"
-            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_TOPIC_ARN))
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(_MESSAGING_DESTINATION))
     elif is_db_span(span):
         remote_resource_type = _DB_CONNECTION_STRING_TYPE
         remote_resource_identifier = _get_db_connection(span)
diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
index 717d5f58e..9730f2d2a 100644
--- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
+++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py
@@ -4,22 +4,19 @@
 import importlib
 
 from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS
-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
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace.span import Span
 
 
 def _apply_botocore_instrumentation_patches() -> None:
     """Botocore instrumentation patches
 
-    Adds patches to provide additional support for Kinesis, S3, SQS and SNS.
+    Adds patches to provide additional support for Kinesis, S3, and SQS.
     """
     _apply_botocore_kinesis_patch()
     _apply_botocore_s3_patch()
     _apply_botocore_sqs_patch()
-    _apply_botocore_sns_patch()
 
 
 def _apply_botocore_kinesis_patch() -> None:
@@ -68,36 +65,6 @@ def patch_extract_attributes(self, attributes: _AttributeMapT):
     _SqsExtension.extract_attributes = patch_extract_attributes
 
 
-def _apply_botocore_sns_patch() -> None:
-    """Botocore instrumentation patch for SNS
-
-    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.
-    """
-    old_extract_attributes = _SnsExtension.extract_attributes
-
-    def patch_extract_attributes(self, attributes: _AttributeMapT):
-        old_extract_attributes(self, attributes)
-        topic_arn = self._call_context.params.get("TopicArn")
-        if topic_arn:
-            attributes["aws.sns.topic_arn"] = topic_arn
-
-    old_on_success = _SnsExtension.on_success
-
-    def patch_on_success(self, span: Span, result: _BotoResultT):
-        old_on_success(self, span, result)
-        topic_arn = result.get("TopicArn")
-        if topic_arn:
-            span.set_attribute(
-                "aws.sns.topic_arn",
-                topic_arn,
-            )
-
-    _SnsExtension.extract_attributes = patch_extract_attributes
-    _SnsExtension.on_success = patch_on_success
-
-
 # The OpenTelemetry Authors code
 def _lazy_load(module, cls):
     """Clone of upstream opentelemetry.instrumentation.botocore.extensions.lazy_load
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py
index 400e0051d..0d9d739cc 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py
+++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py
@@ -19,7 +19,6 @@
     AWS_REMOTE_SERVICE,
     AWS_SPAN_KIND,
     AWS_STREAM_NAME,
-    AWS_TOPIC_ARN,
 )
 from amazon.opentelemetry.distro._aws_metric_attribute_generator import _AwsMetricAttributeGenerator
 from amazon.opentelemetry.distro.metric_attribute_generator import DEPENDENCY_METRIC, SERVICE_METRIC
@@ -980,9 +979,14 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         self._mock_attribute([SpanAttributes.AWS_DYNAMODB_TABLE_NAMES], [None])
 
         # Validate behaviour of AWS_TOPIC_ARN attribute, then remove it
-        self._mock_attribute([AWS_TOPIC_ARN], ["arn:aws:sns:us-west-2:012345678901:test_topic"], keys, values)
+        self._mock_attribute(
+            [SpanAttributes.MESSAGING_DESTINATION, SpanAttributes.RPC_SERVICE],
+            ["arn:aws:sns:us-west-2:012345678901:test_topic", "SNS"],
+            keys,
+            values,
+        )
         self._validate_remote_resource_attributes("AWS::SNS::Topic", "arn:aws:sns:us-west-2:012345678901:test_topic")
-        self._mock_attribute([AWS_TOPIC_ARN], [None])
+        self._mock_attribute([SpanAttributes.MESSAGING_DESTINATION, SpanAttributes.RPC_SERVICE], [None, None])
 
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
 
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
index 80a8a88c5..3a5641d1a 100644
--- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py
+++ b/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 Any, Dict
+from typing import Dict
 from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
@@ -9,13 +9,11 @@
 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"
 _QUEUE_NAME: str = "queueName"
 _QUEUE_URL: str = "queueUrl"
-_TOPIC_ARN: str = "topicArn"
 
 # Patch names
 GET_DISTRIBUTION_PATCH: str = (
@@ -95,13 +93,6 @@ def _test_unpatched_botocore_instrumentation(self):
         self.assertFalse("aws.sqs.queue_url" in sqs_attributes)
         self.assertFalse("aws.sqs.queue_name" in sqs_attributes)
 
-        # SNS
-        self.assertTrue("sns" in _KNOWN_EXTENSIONS, "Upstream has removed the SNS extension")
-        sns_attributes: Dict[str, str] = _do_extract_sns_attributes()
-        self.assertFalse("aws.sns.topic_arn" in sns_attributes)
-        sns_success_attributes: Dict[str, str] = _do_sns_on_success()
-        self.assertFalse("aws.sns.topic_arn" in sns_success_attributes)
-
     def _test_patched_botocore_instrumentation(self):
         # Kinesis
         self.assertTrue("kinesis" in _KNOWN_EXTENSIONS)
@@ -124,16 +115,6 @@ def _test_patched_botocore_instrumentation(self):
         self.assertTrue("aws.sqs.queue_name" in sqs_attributes)
         self.assertEqual(sqs_attributes["aws.sqs.queue_name"], _QUEUE_NAME)
 
-        # SNS
-        self.assertTrue("sns" in _KNOWN_EXTENSIONS)
-        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 _test_botocore_installed_flag(self):
         with patch(
             "amazon.opentelemetry.distro.patches._botocore_patches._apply_botocore_instrumentation_patches"
@@ -175,18 +156,6 @@ def _do_extract_sqs_attributes() -> Dict[str, str]:
     return _do_extract_attributes(service_name, params)
 
 
-def _do_extract_sns_attributes() -> Dict[str, str]:
-    service_name: str = "sns"
-    params: Dict[str, str] = {"TopicArn": _TOPIC_ARN}
-    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
@@ -194,17 +163,3 @@ def _do_extract_attributes(service_name: str, params: Dict[str, str]) -> Dict[st
     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
diff --git a/contract-tests/images/applications/botocore/botocore_server.py b/contract-tests/images/applications/botocore/botocore_server.py
index 42db6a2e3..cddb943fd 100644
--- a/contract-tests/images/applications/botocore/botocore_server.py
+++ b/contract-tests/images/applications/botocore/botocore_server.py
@@ -48,7 +48,7 @@ def do_GET(self):
 
     # pylint: disable=invalid-name
     def do_POST(self):
-        if self.in_path("sqserror") or self.in_path("snserror"):
+        if self.in_path("sqserror"):
             self.send_response(self.main_status)
             self.send_header("Content-type", "text/xml")
             self.end_headers()
@@ -205,47 +205,6 @@ def _handle_kinesis_request(self) -> None:
         else:
             set_main_status(404)
 
-    def _handle_sns_request(self) -> None:
-        sns_client: BaseClient = boto3.client("sns", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
-        if self.in_path(_ERROR):
-            set_main_status(400)
-            try:
-                error_client: BaseClient = boto3.client(
-                    "sns", endpoint_url=_ERROR_ENDPOINT + "/snserror", region_name=_AWS_REGION
-                )
-                topic_arn = "arn:aws:sns:us-west-2:000000000000:test_topic/snserror"
-                message = "Hello from Amazon SNS!"
-                subject = "Test Message"
-                message_attributes = {"Attribute1": {"DataType": "String", "StringValue": "Value1"}}
-                error_client.publish(
-                    TopicArn=topic_arn, Message=message, Subject=subject, MessageAttributes=message_attributes
-                )
-            except Exception as exception:
-                print("Expected exception occurred", exception)
-        elif self.in_path(_FAULT):
-            set_main_status(500)
-            try:
-                fault_client: BaseClient = boto3.client(
-                    "sns", endpoint_url=_FAULT_ENDPOINT, region_name=_AWS_REGION, config=_NO_RETRY_CONFIG
-                )
-                fault_client.get_topic_attributes(TopicArn="invalid_topic_arn")
-            except Exception as exception:
-                print("Expected exception occurred", exception)
-        elif self.in_path("gettopattributes/get-topic-attributes"):
-            set_main_status(200)
-            sns_client.get_topic_attributes(TopicArn="arn:aws:sns:us-west-2:000000000000:test_topic")
-        elif self.in_path("publishmessage/publish-message/some-message"):
-            set_main_status(200)
-            topic_arn = "arn:aws:sns:us-west-2:000000000000:test_topic"
-            message = "Hello from Amazon SNS!"
-            subject = "Test Message"
-            message_attributes = {"Attribute1": {"DataType": "String", "StringValue": "Value1"}}
-            sns_client.publish(
-                TopicArn=topic_arn, Message=message, Subject=subject, MessageAttributes=message_attributes
-            )
-        else:
-            set_main_status(404)
-
     def _end_request(self, status_code: int):
         self.send_response_only(status_code)
         self.end_headers()
@@ -290,10 +249,6 @@ def prepare_aws_server() -> None:
         # Set up Kinesis so tests can access a stream.
         kinesis_client: BaseClient = boto3.client("kinesis", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
         kinesis_client.create_stream(StreamName="test_stream", ShardCount=1)
-
-        # Set up SNS so tests can access a topic.
-        sns_client: BaseClient = boto3.client("sns", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
-        sns_client.create_topic(Name="test_topic")
     except Exception as exception:
         print("Unexpected exception occurred", exception)
 
diff --git a/contract-tests/tests/test/amazon/botocore/botocore_test.py b/contract-tests/tests/test/amazon/botocore/botocore_test.py
index 2012446da..489885faf 100644
--- a/contract-tests/tests/test/amazon/botocore/botocore_test.py
+++ b/contract-tests/tests/test/amazon/botocore/botocore_test.py
@@ -29,7 +29,6 @@
 _AWS_QUEUE_URL: str = "aws.sqs.queue_url"
 _AWS_QUEUE_NAME: str = "aws.sqs.queue_name"
 _AWS_STREAM_NAME: str = "aws.kinesis.stream_name"
-_AWS_TOPIC_ARN: str = "aws.sns.topic_arn"
 
 
 # pylint: disable=too-many-public-methods
@@ -390,78 +389,6 @@ def test_kinesis_fault(self):
             span_kind="CLIENT",
         )
 
-    def test_sns_get_topic_attributes(self):
-        self.do_test_requests(
-            "sns/gettopattributes/get-topic-attributes",
-            "GET",
-            200,
-            0,
-            0,
-            remote_service="AWS::SNS",
-            remote_operation="GetTopicAttributes",
-            remote_resource_type="AWS::SNS::Topic",
-            remote_resource_identifier="arn:aws:sns:us-west-2:000000000000:test_topic",
-            request_specific_attributes={
-                _AWS_TOPIC_ARN: "arn:aws:sns:us-west-2:000000000000:test_topic",
-            },
-            span_name="SNS.GetTopicAttributes",
-            span_kind="CLIENT",
-        )
-
-    def test_sns_publish_message(self):
-        self.do_test_requests(
-            "sns/publishmessage/publish-message/some-message",
-            "GET",
-            200,
-            0,
-            0,
-            remote_service="AWS::SNS",
-            remote_operation="Publish",
-            remote_resource_type="AWS::SNS::Topic",
-            remote_resource_identifier="arn:aws:sns:us-west-2:000000000000:test_topic",
-            request_specific_attributes={
-                _AWS_TOPIC_ARN: "arn:aws:sns:us-west-2:000000000000:test_topic",
-            },
-            span_name="test_topic send",
-            span_kind="PRODUCER",
-        )
-
-    def test_sns_error(self):
-        self.do_test_requests(
-            "sns/error",
-            "GET",
-            400,
-            1,
-            0,
-            remote_service="AWS::SNS",
-            remote_operation="Publish",
-            remote_resource_type="AWS::SNS::Topic",
-            remote_resource_identifier="arn:aws:sns:us-west-2:000000000000:test_topic/snserror",
-            request_specific_attributes={
-                _AWS_TOPIC_ARN: "arn:aws:sns:us-west-2:000000000000:test_topic/snserror",
-            },
-            span_name="test_topic/snserror send",
-            span_kind="PRODUCER",
-        )
-
-    def test_sns_fault(self):
-        self.do_test_requests(
-            "sns/fault",
-            "GET",
-            500,
-            0,
-            1,
-            remote_service="AWS::SNS",
-            remote_operation="GetTopicAttributes",
-            remote_resource_type="AWS::SNS::Topic",
-            remote_resource_identifier="invalid_topic_arn",
-            request_specific_attributes={
-                _AWS_TOPIC_ARN: "invalid_topic_arn",
-            },
-            span_name="SNS.GetTopicAttributes",
-            span_kind="CLIENT",
-        )
-
     @override
     def _assert_aws_span_attributes(self, resource_scope_spans: List[ResourceScopeSpan], path: str, **kwargs) -> None:
         target_spans: List[Span] = []

From 602ce008a559222d1c797b1a2468ba60b291c24b Mon Sep 17 00:00:00 2001
From: Zhonghao Zhao <zzhlogin@amazon.com>
Date: Tue, 25 Jun 2024 12:29:07 -0700
Subject: [PATCH 16/17] Apply lint.

---
 contract-tests/images/applications/botocore/botocore_server.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/contract-tests/images/applications/botocore/botocore_server.py b/contract-tests/images/applications/botocore/botocore_server.py
index cddb943fd..dd1e34c6b 100644
--- a/contract-tests/images/applications/botocore/botocore_server.py
+++ b/contract-tests/images/applications/botocore/botocore_server.py
@@ -41,8 +41,6 @@ def do_GET(self):
             self._handle_sqs_request()
         if self.in_path("kinesis"):
             self._handle_kinesis_request()
-        if self.in_path("sns"):
-            self._handle_sns_request()
 
         self._end_request(self.main_status)
 

From 3840d90315052b721b632a7dd33c50dd7c73b629 Mon Sep 17 00:00:00 2001
From: Thomas Pierce <thp@amazon.com>
Date: Fri, 23 Aug 2024 16:45:55 -0700
Subject: [PATCH 17/17] Update Dockerfile

---
 Dockerfile | 41 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 37 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index e8aa35db1..23fc2c907 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,7 +3,9 @@
 #   The packages are installed in the `/autoinstrumentation` directory. This is required as when instrumenting the pod by CWOperator,
 #   one init container will be created to copy all the content in `/autoinstrumentation` directory to app's container. Then
 #   update the `PYTHONPATH` environment variable accordingly. Then in the second stage, copy the directory to `/autoinstrumentation`.
-FROM python:3.11 AS build
+
+# Stage 1: Install ADOT Python in the /operator-build folder
+FROM public.ecr.aws/docker/library/python:3.11 AS build
 
 WORKDIR /operator-build
 
@@ -18,11 +20,42 @@ RUN sed -i "/opentelemetry-exporter-otlp-proto-grpc/d" ./aws-opentelemetry-distr
 
 RUN mkdir workspace && pip install --target workspace ./aws-opentelemetry-distro
 
-FROM public.ecr.aws/amazonlinux/amazonlinux:minimal
+# Stage 2: Build the cp-utility binary
+FROM public.ecr.aws/docker/library/rust:1.75 as builder
+
+WORKDIR /usr/src/cp-utility
+COPY ./tools/cp-utility .
+
+## TARGETARCH is defined by buildx
+# https://docs.docker.com/engine/reference/builder/#automatic-platform-args-in-the-global-scope
+ARG TARGETARCH
+
+# Run validations and audit only on amd64 because it is faster and those two steps
+# are only used to validate the source code and don't require anything that is
+# architecture specific.
+
+# Validations
+# Validate formatting
+RUN if [ $TARGETARCH = "amd64" ]; then rustup component add rustfmt && cargo fmt --check ; fi
+
+# Audit dependencies
+RUN if [ $TARGETARCH = "amd64" ]; then cargo install cargo-audit && cargo audit ; fi
+
+
+# Cross-compile based on the target platform.
+RUN if [ $TARGETARCH = "amd64" ]; then export ARCH="x86_64" ; \
+    elif [ $TARGETARCH = "arm64" ]; then export ARCH="aarch64" ; \
+    else false; \
+    fi \
+    && rustup target add ${ARCH}-unknown-linux-musl \
+    && cargo test  --target ${ARCH}-unknown-linux-musl \
+    && cargo install --target ${ARCH}-unknown-linux-musl --path . --root .
+
+# Stage 3: Build the distribution image by copying the THIRD-PARTY-LICENSES, the custom built cp command from stage 2, and the installed ADOT Python from stage 1 to their respective destinations
+FROM scratch
 
 # Required to copy attribute files to distributed docker images
 ADD THIRD-PARTY-LICENSES ./THIRD-PARTY-LICENSES
 
+COPY --from=builder /usr/src/cp-utility/bin/cp-utility /bin/cp
 COPY --from=build /operator-build/workspace /autoinstrumentation
-
-RUN chmod -R go+r /autoinstrumentation