Add SecretsManager Secret Arn to Span Data

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py

@@ -17,8 +17,10 @@
 # TODO：Move to Semantic Conventions when these attributes are added.
 AWS_SQS_QUEUE_URL: str = "aws.sqs.queue.url"
 AWS_SQS_QUEUE_NAME: str = "aws.sqs.queue.name"
+AWS_SNS_TOPIC_ARN: str = "aws.sns.topic.arn"
 AWS_KINESIS_STREAM_NAME: str = "aws.kinesis.stream.name"
 AWS_BEDROCK_DATA_SOURCE_ID: str = "aws.bedrock.data_source.id"
 AWS_BEDROCK_KNOWLEDGE_BASE_ID: str = "aws.bedrock.knowledge_base.id"
 AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id"
 AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail.id"
+AWS_SECRETSMANAGER_SECRET_ARN: str = "aws.secretsmanager.secret.arn"

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py

@@ -18,6 +18,8 @@
     AWS_REMOTE_RESOURCE_IDENTIFIER,
     AWS_REMOTE_RESOURCE_TYPE,
     AWS_REMOTE_SERVICE,
+    AWS_SECRETSMANAGER_SECRET_ARN,
+    AWS_SNS_TOPIC_ARN,
     AWS_SPAN_KIND,
     AWS_SQS_QUEUE_NAME,
     AWS_SQS_QUEUE_URL,
@@ -85,8 +87,10 @@
 _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"
 _NORMALIZED_BEDROCK_SERVICE_NAME: str = "AWS::Bedrock"
 _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME: str = "AWS::BedrockRuntime"
+_NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager"
 _DB_CONNECTION_STRING_TYPE: str = "DB::Connection"
 
 # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
@@ -308,6 +312,7 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str
             "Bedrock Agent": _NORMALIZED_BEDROCK_SERVICE_NAME,
             "Bedrock Agent Runtime": _NORMALIZED_BEDROCK_SERVICE_NAME,
             "Bedrock Runtime": _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME,
+            "Secrets Manager": _NORMALIZED_SECRETSMANAGER_SERVICE_NAME,
         }
         return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name)
     return service_name
@@ -407,6 +412,12 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
         elif is_key_present(span, GEN_AI_REQUEST_MODEL):
             remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Model"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(GEN_AI_REQUEST_MODEL))
+        elif is_key_present(span, AWS_SNS_TOPIC_ARN):
+            remote_resource_type = _NORMALIZED_SNS_SERVICE_NAME + "::Topic"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_SNS_TOPIC_ARN))
+        elif is_key_present(span, AWS_SECRETSMANAGER_SECRET_ARN):
+            remote_resource_type = _NORMALIZED_SECRETSMANAGER_SERVICE_NAME + "::Secret"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_SECRETSMANAGER_SECRET_ARN))
     elif is_db_span(span):
         remote_resource_type = _DB_CONNECTION_STRING_TYPE
         remote_resource_identifier = _get_db_connection(span)

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

@@ -15,9 +15,11 @@
     _BedrockRuntimeExtension,
 )
 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:
@@ -29,6 +31,41 @@ def _apply_botocore_instrumentation_patches() -> None:
     _apply_botocore_s3_patch()
     _apply_botocore_sqs_patch()
     _apply_botocore_bedrock_patch()
+    _apply_botocore_sns_patch()
+    _apply_botocore_secretsmanager_patch()
+
+
+def _apply_botocore_secretsmanager_patch() -> None:
+    """Botocore instrumentation patch for SecretsManager
+
+    This patch adds an extension to the upstream's list of known extensions for SecretsManager.
+    Extensions allow for custom logic for adding service-specific information to spans,
+    such as attributes. Specifically, we are adding logic to add the
+    AWS_SECRETSMANAGER_SECRET_ARN attribute.
+    """
+    _KNOWN_EXTENSIONS["secretsmanager"] = _lazy_load(".", "_SecretsManagerExtension")
+
+
+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 the `aws.sns.topic.arn` attribute,
+    to be used to generate AWS_REMOTE_RESOURCE_TYPE and AWS_REMOTE_RESOURCE_IDENTIFIER.
+    There exists SpanAttributes.MESSAGING_DESTINATION_NAME in the upstream logic that we could
+    re-purpose here. However, we are not using it here to maintain consistent naming patterns
+    with other AWS resources.
+    """
+    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
 
 
 def _apply_botocore_kinesis_patch() -> None:
@@ -108,6 +145,24 @@ def loader():
 # END The OpenTelemetry Authors code
 
 
+class _SecretsManagerExtension(_AwsSdkExtension):
+    def extract_attributes(self, attributes: _AttributeMapT):
+        """
+        SecretId can be secret name or secret arn, the function extracts attributes only if
+        the SecretId parameter is provided as an arn which starts with
+        'arn:aws:secretsmanager:'.
+        """
+        secret_id = self._call_context.params.get("SecretId")
+        if secret_id and secret_id.startswith("arn:aws:secretsmanager:"):
+            attributes["aws.secretsmanager.secret.arn"] = secret_id
+
+    # pylint: disable=no-self-use
+    def on_success(self, span: Span, result: _BotoResultT):
+        secret_arn = result.get("ARN")
+        if secret_arn:
+            span.set_attribute("aws.secretsmanager.secret.arn", secret_arn)
+
+
 class _S3Extension(_AwsSdkExtension):
     def extract_attributes(self, attributes: _AttributeMapT):
         bucket_name = self._call_context.params.get("Bucket")

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

@@ -21,6 +21,8 @@
     AWS_REMOTE_RESOURCE_IDENTIFIER,
     AWS_REMOTE_RESOURCE_TYPE,
     AWS_REMOTE_SERVICE,
+    AWS_SECRETSMANAGER_SECRET_ARN,
+    AWS_SNS_TOPIC_ARN,
     AWS_SPAN_KIND,
     AWS_SQS_QUEUE_NAME,
     AWS_SQS_QUEUE_URL,
@@ -877,6 +879,8 @@ def test_normalize_remote_service_name_aws_sdk(self):
         self.validate_aws_sdk_service_normalization("Bedrock Agent", "AWS::Bedrock")
         self.validate_aws_sdk_service_normalization("Bedrock Agent Runtime", "AWS::Bedrock")
         self.validate_aws_sdk_service_normalization("Bedrock Runtime", "AWS::BedrockRuntime")
+        self.validate_aws_sdk_service_normalization("SNS", "AWS::SNS")
+        self.validate_aws_sdk_service_normalization("Secrets Manager", "AWS::SecretsManager")
 
     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])
@@ -1001,6 +1005,11 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         self._validate_remote_resource_attributes("AWS::SQS::Queue", "aws_queue_name")
         self._mock_attribute([AWS_SQS_QUEUE_URL, AWS_SQS_QUEUE_NAME], [None, None])
 
+        # Validate behaviour of AWS_SNS_TOPIC_ARN attribute, then remove it.
+        self._mock_attribute([AWS_SNS_TOPIC_ARN], ["arn:aws:sns:us-west-2:012345678901:test_topic"], keys, values)
+        self._validate_remote_resource_attributes("AWS::SNS::Topic", "arn:aws:sns:us-west-2:012345678901:test_topic")
+        self._mock_attribute([AWS_SNS_TOPIC_ARN], [None])
+
         # Validate behaviour of AWS_KINESIS_STREAM_NAME attribute, then remove it.
         self._mock_attribute([AWS_KINESIS_STREAM_NAME], ["aws_stream_name"], keys, values)
         self._validate_remote_resource_attributes("AWS::Kinesis::Stream", "aws_stream_name")
@@ -1083,6 +1092,18 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         self._validate_remote_resource_attributes("AWS::Bedrock::Model", "test.service_^^id")
         self._mock_attribute([GEN_AI_REQUEST_MODEL], [None])
 
+        # Validate behaviour of AWS_SECRETSMANAGER_SECRET_ARN attribute, then remove it
+        self._mock_attribute(
+            [AWS_SECRETSMANAGER_SECRET_ARN],
+            ["arn:aws:secretsmanager:us-east-1:123456789012:secret:secret_name-lERW9H"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes(
+            "AWS::SecretsManager::Secret", "arn:aws:secretsmanager:us-east-1:123456789012:secret:secret_name-lERW9H"
+        )
+        self._mock_attribute([AWS_SECRETSMANAGER_SECRET_ARN], [None])
+
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
 
     def test_client_db_span_with_remote_resource_attributes(self):

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

@@ -20,12 +20,14 @@
 _BUCKET_NAME: str = "bucketName"
 _QUEUE_NAME: str = "queueName"
 _QUEUE_URL: str = "queueUrl"
+_TOPIC_ARN: str = "topicArn"
 _BEDROCK_AGENT_ID: str = "agentId"
 _BEDROCK_DATASOURCE_ID: str = "DataSourceId"
 _BEDROCK_GUARDRAIL_ID: str = "GuardrailId"
 _BEDROCK_KNOWLEDGEBASE_ID: str = "KnowledgeBaseId"
 _GEN_AI_SYSTEM: str = "aws_bedrock"
 _GEN_AI_REQUEST_MODEL: str = "genAiReuqestModelId"
+_SECRET_ARN: str = "arn:aws:secretsmanager:us-west-2:000000000000:secret:testSecret-ABCDEF"
 
 # Patch names
 GET_DISTRIBUTION_PATCH: str = (
@@ -127,6 +129,11 @@ def _test_unpatched_botocore_instrumentation(self):
         self.assertFalse("aws.sqs.queue.url" in attributes)
         self.assertFalse("aws.sqs.queue.name" in 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)
+
         # Bedrock
         self.assertFalse("bedrock" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock extension")
 
@@ -141,6 +148,9 @@ def _test_unpatched_botocore_instrumentation(self):
         # BedrockRuntime
         self.assertFalse("bedrock-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a bedrock-runtime extension")
 
+        # SecretsManager
+        self.assertFalse("secretsmanager" in _KNOWN_EXTENSIONS, "Upstream has added a SecretsManager extension")
+
     def _test_unpatched_gevent_instrumentation(self):
         self.assertFalse(gevent.monkey.is_module_patched("os"), "gevent os module has been patched")
         self.assertFalse(gevent.monkey.is_module_patched("thread"), "gevent thread module has been patched")
@@ -178,6 +188,12 @@ 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)
+
         # Bedrock
         self._test_patched_bedrock_instrumentation()
 
@@ -200,6 +216,15 @@ def _test_patched_botocore_instrumentation(self):
         self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM)
         self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], _GEN_AI_REQUEST_MODEL)
 
+        # SecretsManager
+        self.assertTrue("secretsmanager" in _KNOWN_EXTENSIONS)
+        secretsmanager_attributes: Dict[str, str] = _do_extract_secretsmanager_attributes()
+        self.assertTrue("aws.secretsmanager.secret.arn" in secretsmanager_attributes)
+        self.assertEqual(secretsmanager_attributes["aws.secretsmanager.secret.arn"], _SECRET_ARN)
+        secretsmanager_success_attributes: Dict[str, str] = _do_on_success_secretsmanager()
+        self.assertTrue("aws.secretsmanager.secret.arn" in secretsmanager_success_attributes)
+        self.assertEqual(secretsmanager_success_attributes["aws.secretsmanager.secret.arn"], _SECRET_ARN)
+
     def _test_patched_gevent_os_ssl_instrumentation(self):
         # Only ssl and os module should have been patched since the environment variable was set to 'os, ssl'
         self.assertTrue(gevent.monkey.is_module_patched("ssl"), "gevent ssl module has not been patched")
@@ -325,6 +350,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_bedrock(service, operation=None) -> Dict[str, str]:
     params: Dict[str, Any] = {
         "agentId": _BEDROCK_AGENT_ID,
@@ -347,6 +378,18 @@ def _do_on_success_bedrock(service, operation=None) -> Dict[str, str]:
     return _do_on_success(service, result, operation)
 
 
+def _do_extract_secretsmanager_attributes() -> Dict[str, str]:
+    service_name: str = "secretsmanager"
+    params: Dict[str, str] = {"SecretId": _SECRET_ARN}
+    return _do_extract_attributes(service_name, params)
+
+
+def _do_on_success_secretsmanager() -> Dict[str, str]:
+    service_name: str = "secretsmanager"
+    result: Dict[str, Any] = {"ARN": _SECRET_ARN}
+    return _do_on_success(service_name, result)
+
+
 def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]:
     mock_call_context: MagicMock = MagicMock()
     mock_call_context.params = params