feat: Support a list of new AWS resources for AWS Python SDK

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

@@ -23,3 +23,10 @@
 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_SNS_TOPIC_ARN: str = "aws.sns.topic.arn"
+AWS_STEPFUNCTIONS_STATEMACHINE_ARN: str = "aws.stepfunctions.state_machine.arn"
+AWS_STEPFUNCTIONS_ACTIVITY_ARN: str = "aws.stepfunctions.activity.arn"
+AWS_LAMBDA_FUNCTION_NAME: str = "aws.lambda.function.name"
+AWS_LAMBDA_RESOURCEMAPPING_ID: str = "aws.lambda.resource_mapping.id"
+AWS_LAMBDA_FUNCTION_ARN: str = "aws.lambda.function.arn"

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

@@ -12,16 +12,23 @@
     AWS_BEDROCK_KNOWLEDGE_BASE_ID,
     AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
     AWS_KINESIS_STREAM_NAME,
+    AWS_LAMBDA_FUNCTION_ARN,
+    AWS_LAMBDA_FUNCTION_NAME,
+    AWS_LAMBDA_RESOURCEMAPPING_ID,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
     AWS_REMOTE_DB_USER,
     AWS_REMOTE_OPERATION,
     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,
+    AWS_STEPFUNCTIONS_ACTIVITY_ARN,
+    AWS_STEPFUNCTIONS_STATEMACHINE_ARN,
 )
 from amazon.opentelemetry.distro._aws_span_processing_util import (
     GEN_AI_REQUEST_MODEL,
@@ -88,6 +95,10 @@
 _NORMALIZED_SQS_SERVICE_NAME: str = "AWS::SQS"
 _NORMALIZED_BEDROCK_SERVICE_NAME: str = "AWS::Bedrock"
 _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME: str = "AWS::BedrockRuntime"
+_NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager"
+_NORMALIZED_SNS_SERVICE_NAME: str = "AWS::SNS"
+_NORMALIZED_STEPFUNCTIONS_SERVICE_NAME: str = "AWS::StepFunctions"
+_NORMALIZED_LAMBDA_SERVICE_NAME: str = "AWS::Lambda"
 _DB_CONNECTION_STRING_TYPE: str = "DB::Connection"
 
 # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
@@ -309,6 +320,9 @@ 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,
+            "SNS": _NORMALIZED_SNS_SERVICE_NAME,
+            "SFN": _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME,
         }
         return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name)
     return service_name
@@ -359,7 +373,7 @@ def _generate_remote_operation(span: ReadableSpan) -> str:
     return remote_operation
 
 
-# pylint: disable=too-many-branches
+# pylint: disable=too-many-branches,too-many-statements
 def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttributes) -> None:
     """
     Remote resource attributes {@link AwsAttributeKeys#AWS_REMOTE_RESOURCE_TYPE} and {@link
@@ -416,6 +430,37 @@ 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_SECRETSMANAGER_SECRET_ARN):
+            remote_resource_type = _NORMALIZED_SECRETSMANAGER_SERVICE_NAME + "::Secret"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_SECRETSMANAGER_SECRET_ARN)).split(
+                ":"
+            )[-1]
+            cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_SECRETSMANAGER_SECRET_ARN))
+        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)).split(":")[-1]
+            cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_SNS_TOPIC_ARN))
+        elif is_key_present(span, AWS_STEPFUNCTIONS_STATEMACHINE_ARN):
+            remote_resource_type = _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME + "::StateMachine"
+            remote_resource_identifier = _escape_delimiters(
+                span.attributes.get(AWS_STEPFUNCTIONS_STATEMACHINE_ARN)
+            ).split(":")[-1]
+            cloudformation_primary_identifier = _escape_delimiters(
+                span.attributes.get(AWS_STEPFUNCTIONS_STATEMACHINE_ARN)
+            )
+        elif is_key_present(span, AWS_STEPFUNCTIONS_ACTIVITY_ARN):
+            remote_resource_type = _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME + "::Activity"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN)).split(
+                ":"
+            )[-1]
+            cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN))
+        elif is_key_present(span, AWS_LAMBDA_RESOURCEMAPPING_ID):
+            remote_resource_type = _NORMALIZED_LAMBDA_SERVICE_NAME + "::EventSourceMapping"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_RESOURCEMAPPING_ID))
+        elif is_key_present(span, AWS_LAMBDA_FUNCTION_NAME):
+            remote_resource_type = _NORMALIZED_LAMBDA_SERVICE_NAME + "::Function"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_FUNCTION_NAME))
+            cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_FUNCTION_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

@@ -5,8 +5,15 @@
 
 from amazon.opentelemetry.distro._aws_attribute_keys import (
     AWS_KINESIS_STREAM_NAME,
+    AWS_LAMBDA_FUNCTION_ARN,
+    AWS_LAMBDA_FUNCTION_NAME,
+    AWS_LAMBDA_RESOURCEMAPPING_ID,
+    AWS_SECRETSMANAGER_SECRET_ARN,
+    AWS_SNS_TOPIC_ARN,
     AWS_SQS_QUEUE_NAME,
     AWS_SQS_QUEUE_URL,
+    AWS_STEPFUNCTIONS_ACTIVITY_ARN,
+    AWS_STEPFUNCTIONS_STATEMACHINE_ARN,
 )
 from amazon.opentelemetry.distro.patches._bedrock_patches import (  # noqa # pylint: disable=unused-import
     _BedrockAgentExtension,
@@ -15,9 +22,12 @@
     _BedrockRuntimeExtension,
 )
 from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS
+from opentelemetry.instrumentation.botocore.extensions.lmbd import _LambdaExtension
+from opentelemetry.instrumentation.botocore.extensions.sns import _SnsExtension
 from opentelemetry.instrumentation.botocore.extensions.sqs import _SqsExtension
-from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension
+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 +39,100 @@ def _apply_botocore_instrumentation_patches() -> None:
     _apply_botocore_s3_patch()
     _apply_botocore_sqs_patch()
     _apply_botocore_bedrock_patch()
+    _apply_botocore_secretsmanager_patch()
+    _apply_botocore_sns_patch()
+    _apply_botocore_stepfunctions_patch()
+    _apply_botocore_lambda_patch()
+
+
+def _apply_botocore_lambda_patch() -> None:
+    """Botocore instrumentation patch for Lambda
+
+    This patch adds an extension to the upstream's list of known extensions for Lambda.
+    Extensions allow for custom logic for adding service-specific information to spans,
+    such as attributes. Specifically, we are adding logic to add the
+    `aws.lambda.function.name` and  `aws.lambda.resource_mapping.id` attributes
+
+    Sidenote: There exists SpanAttributes.FAAS_INVOKED_NAME for invoke operations
+    in upstream. However, we want to cover more operations to extract 'FunctionName',
+    so we define `aws.lambda.function.name` separately. Additionally, this helps
+    us maintain naming consistency with the other AWS resources.
+    """
+    old_extract_attributes = _LambdaExtension.extract_attributes
+
+    def patch_extract_attributes(self, attributes: _AttributeMapT):
+        old_extract_attributes(self, attributes)
+        # This param can be passed as an arn or a name. We standardize it to be the name.
+        function_name_param = self._call_context.params.get("FunctionName")
+        if function_name_param:
+            function_name = function_name_param
+            if function_name_param.startswith("arn:aws:lambda:"):
+                function_name = function_name_param.split(":")[-1]
+            attributes[AWS_LAMBDA_FUNCTION_NAME] = function_name
+        resource_mapping_id = self._call_context.params.get("UUID")
+        if resource_mapping_id:
+            attributes[AWS_LAMBDA_RESOURCEMAPPING_ID] = resource_mapping_id
+
+    old_on_success = _LambdaExtension.on_success
+
+    def patch_on_success(self, span: Span, result: _BotoResultT):
+        old_on_success(self, span, result)
+        lambda_configuration = result.get("Configuration", {})
+        function_arn = lambda_configuration.get("FunctionArn")
+        if function_arn:
+            span.set_attribute(AWS_LAMBDA_FUNCTION_ARN, function_arn)
+
+    _LambdaExtension.extract_attributes = patch_extract_attributes
+    _LambdaExtension.on_success = patch_on_success
+
+
+def _apply_botocore_stepfunctions_patch() -> None:
+    """Botocore instrumentation patch for StepFunctions
+
+    This patch adds an extension to the upstream's list of known extensions for
+    StepFunctions. Extensions allow for custom logic for adding service-specific
+    information to spans, such as attributes. Specifically, we are adding logic
+    to add the `aws.stepfunctions.state_machine.arn` and `aws.stepfunctions.activity.arn`
+    attributes, to be used to generate RemoteTarget and achieve partity with the
+    Java instrumentation.
+    """
+    _KNOWN_EXTENSIONS["stepfunctions"] = _lazy_load(".", "_StepFunctionsExtension")
+
+
+def _apply_botocore_sns_patch() -> None:
+    """Botocore instrumentation patch for SNS
+
+    This patch adds an extension to the upstream's list of known extensions 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 RemoteTarget and achieve
+    parity with the Java instrumentation.
+
+    Sidenote: There exists SpanAttributes.MESSAGING_DESTINATION_NAME in the upstream
+    logic that we could re-purpose here. We do not use 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_secretsmanager_patch() -> None:
+    """Botocore instrumentation patch for SecretsManager
+
+    This patch adds an extension to the upstream's list of known extension 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, to be used to generate RemoteTarget and achieve parity with the Java
+    instrumentation.
+    """
+    _KNOWN_EXTENSIONS["secretsmanager"] = _lazy_load(".", "_SecretsManagerExtension")
 
 
 def _apply_botocore_kinesis_patch() -> None:
@@ -108,6 +212,34 @@ def loader():
 # END The OpenTelemetry Authors code
 
 
+class _StepFunctionsExtension(_AwsSdkExtension):
+    def extract_attributes(self, attributes: _AttributeMapT):
+        state_machine_arn = self._call_context.params.get("stateMachineArn")
+        if state_machine_arn:
+            attributes[AWS_STEPFUNCTIONS_STATEMACHINE_ARN] = state_machine_arn
+        activity_arn = self._call_context.params.get("activityArn")
+        if activity_arn:
+            attributes[AWS_STEPFUNCTIONS_ACTIVITY_ARN] = activity_arn
+
+
+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

@@ -14,16 +14,22 @@
     AWS_BEDROCK_KNOWLEDGE_BASE_ID,
     AWS_CONSUMER_PARENT_SPAN_KIND,
     AWS_KINESIS_STREAM_NAME,
+    AWS_LAMBDA_FUNCTION_NAME,
+    AWS_LAMBDA_RESOURCEMAPPING_ID,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
     AWS_REMOTE_DB_USER,
     AWS_REMOTE_OPERATION,
     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,
+    AWS_STEPFUNCTIONS_ACTIVITY_ARN,
+    AWS_STEPFUNCTIONS_STATEMACHINE_ARN,
 )
 from amazon.opentelemetry.distro._aws_metric_attribute_generator import _AwsMetricAttributeGenerator
 from amazon.opentelemetry.distro._aws_span_processing_util import GEN_AI_REQUEST_MODEL
@@ -877,6 +883,9 @@ 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("Secrets Manager", "AWS::SecretsManager")
+        self.validate_aws_sdk_service_normalization("SNS", "AWS::SNS")
+        self.validate_aws_sdk_service_normalization("SFN", "AWS::StepFunctions")
 
     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])
@@ -1093,6 +1102,62 @@ 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", "secret_name-lERW9H")
+        self._mock_attribute([AWS_SECRETSMANAGER_SECRET_ARN], [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", "test_topic")
+        self._mock_attribute([AWS_SNS_TOPIC_ARN], [None])
+
+        # Validate behaviour of AWS_STEPFUNCTIONS_STATEMACHINE_ARN attribute, then remove it.
+        self._mock_attribute(
+            [AWS_STEPFUNCTIONS_STATEMACHINE_ARN],
+            ["arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes("AWS::StepFunctions::StateMachine", "test_state_machine")
+        self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None])
+
+        # Validate behaviour of AWS_STEPFUNCTIONS_ACTIVITY_ARN attribute, then remove it.
+        self._mock_attribute(
+            [AWS_STEPFUNCTIONS_ACTIVITY_ARN],
+            ["arn:aws:states:us-east-1:007003123456789012:activity:testActivity"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes("AWS::StepFunctions::Activity", "testActivity")
+        self._mock_attribute([AWS_STEPFUNCTIONS_ACTIVITY_ARN], [None])
+
+        # Validate behaviour of AWS_LAMBDA_FUNCTION_NAME attribute, then remove it.
+        self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME], ["aws_lambda_function_name"], keys, values)
+        self._validate_remote_resource_attributes("AWS::Lambda::Function", "aws_lambda_function_name")
+        self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME], [None])
+
+        # Validate behaviour of AWS_LAMBDA_RESOURCEMAPPING_ID attribute, then remove it.
+        self._mock_attribute([AWS_LAMBDA_RESOURCEMAPPING_ID], ["aws_event_source_mapping_id"], keys, values)
+        self._validate_remote_resource_attributes("AWS::Lambda::EventSourceMapping", "aws_event_source_mapping_id")
+        self._mock_attribute([AWS_LAMBDA_RESOURCEMAPPING_ID], [None])
+
+        # Validate behaviour of both AWS_LAMBDA_FUNCTION_NAME and AWS_LAMBDA_RESOURCE_MAPPING_ID,
+        # then remove it.
+        self._mock_attribute(
+            [AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_RESOURCEMAPPING_ID],
+            ["aws_lambda_function_name", "aws_event_source_mapping_id"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes("AWS::Lambda::EventSourceMapping", "aws_event_source_mapping_id")
+        self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_RESOURCEMAPPING_ID], [None, None])
+
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
 
     def test_client_db_span_with_remote_resource_attributes(self):