Generate cross-account metric attributes

Author: blairhyy-amazon

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

@@ -7,12 +7,16 @@
 from urllib.parse import ParseResult, urlparse
 
 from amazon.opentelemetry.distro._aws_attribute_keys import (
+    AWS_AUTH_ACCESS_KEY,
+    AWS_AUTH_REGION,
     AWS_BEDROCK_AGENT_ID,
     AWS_BEDROCK_DATA_SOURCE_ID,
     AWS_BEDROCK_GUARDRAIL_ARN,
     AWS_BEDROCK_GUARDRAIL_ID,
     AWS_BEDROCK_KNOWLEDGE_BASE_ID,
     AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
+    AWS_DYNAMODB_TABLE_ARN,
+    AWS_KINESIS_STREAM_ARN,
     AWS_KINESIS_STREAM_NAME,
     AWS_LAMBDA_FUNCTION_ARN,
     AWS_LAMBDA_FUNCTION_NAME,
@@ -22,7 +26,10 @@
     AWS_REMOTE_DB_USER,
     AWS_REMOTE_ENVIRONMENT,
     AWS_REMOTE_OPERATION,
+    AWS_REMOTE_RESOURCE_ACCESS_KEY,
+    AWS_REMOTE_RESOURCE_ACCOUNT_ID,
     AWS_REMOTE_RESOURCE_IDENTIFIER,
+    AWS_REMOTE_RESOURCE_REGION,
     AWS_REMOTE_RESOURCE_TYPE,
     AWS_REMOTE_SERVICE,
     AWS_SECRETSMANAGER_SECRET_ARN,
@@ -56,6 +63,7 @@
     SERVICE_METRIC,
     MetricAttributeGenerator,
 )
+from amazon.opentelemetry.distro.regional_resource_arn_parser import RegionalResourceArnParser
 from amazon.opentelemetry.distro.sqs_url_parser import SqsUrlParser
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import BoundedAttributes, ReadableSpan
@@ -148,7 +156,11 @@ def _generate_dependency_metric_attributes(span: ReadableSpan, resource: Resourc
     _set_service(resource, span, attributes)
     _set_egress_operation(span, attributes)
     _set_remote_service_and_operation(span, attributes)
-    _set_remote_type_and_identifier(span, attributes)
+    is_remote_identifier_present = _set_remote_type_and_identifier(span, attributes)
+    if is_remote_identifier_present:
+        is_remote_account_id_present = _set_remote_account_id_and_region(span, attributes)
+        if not is_remote_account_id_present:
+            _set_remote_access_key_and_region(span, attributes)
     _set_remote_environment(span, attributes)
     _set_remote_db_user(span, attributes)
     _set_span_kind_for_dependency(span, attributes)
@@ -383,7 +395,7 @@ def _generate_remote_operation(span: ReadableSpan) -> str:
 
 
 # pylint: disable=too-many-branches,too-many-statements
-def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttributes) -> None:
+def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttributes) -> bool:
     """
     Remote resource attributes {@link AwsAttributeKeys#AWS_REMOTE_RESOURCE_TYPE} and {@link
     AwsAttributeKeys#AWS_REMOTE_RESOURCE_IDENTIFIER} are used to store information about the resource associated with
@@ -403,9 +415,19 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
         if is_key_present(span, _AWS_TABLE_NAMES) and len(span.attributes.get(_AWS_TABLE_NAMES)) == 1:
             remote_resource_type = _NORMALIZED_DYNAMO_DB_SERVICE_NAME + "::Table"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(_AWS_TABLE_NAMES)[0])
+        elif is_key_present(span, AWS_DYNAMODB_TABLE_ARN):
+            remote_resource_type = _NORMALIZED_DYNAMO_DB_SERVICE_NAME + "::Table"
+            remote_resource_identifier = (
+                _escape_delimiters(span.attributes.get(AWS_DYNAMODB_TABLE_ARN)).split(":")[-1].replace("table/", "")
+            )
         elif is_key_present(span, AWS_KINESIS_STREAM_NAME):
             remote_resource_type = _NORMALIZED_KINESIS_SERVICE_NAME + "::Stream"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_KINESIS_STREAM_NAME))
+        elif is_key_present(span, AWS_KINESIS_STREAM_ARN):
+            remote_resource_type = _NORMALIZED_KINESIS_SERVICE_NAME + "::Stream"
+            remote_resource_identifier = (
+                _escape_delimiters(span.attributes.get(AWS_KINESIS_STREAM_ARN)).split(":")[-1].replace("stream/", "")
+            )
         elif is_key_present(span, _AWS_BUCKET_NAME):
             remote_resource_type = _NORMALIZED_S3_SERVICE_NAME + "::Bucket"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(_AWS_BUCKET_NAME))
@@ -491,6 +513,48 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
         attributes[AWS_REMOTE_RESOURCE_TYPE] = remote_resource_type
         attributes[AWS_REMOTE_RESOURCE_IDENTIFIER] = remote_resource_identifier
         attributes[AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER] = cloudformation_primary_identifier
+        return True
+    return False
+
+
+def _set_remote_account_id_and_region(span: ReadableSpan, attributes: BoundedAttributes) -> bool:
+    ARN_ATTRIBUTES = [
+        AWS_DYNAMODB_TABLE_ARN,
+        AWS_KINESIS_STREAM_ARN,
+        AWS_SNS_TOPIC_ARN,
+        AWS_SECRETSMANAGER_SECRET_ARN,
+        AWS_STEPFUNCTIONS_STATEMACHINE_ARN,
+        AWS_STEPFUNCTIONS_ACTIVITY_ARN,
+        AWS_BEDROCK_GUARDRAIL_ARN,
+        AWS_LAMBDA_FUNCTION_ARN,
+    ]
+    remote_account_id: Optional[str] = None
+    remote_region: Optional[str] = None
+
+    if is_key_present(span, AWS_SQS_QUEUE_URL):
+        queue_url = _escape_delimiters(span.attributes.get(AWS_SQS_QUEUE_URL))
+        remote_account_id = SqsUrlParser.get_account_id(queue_url)
+        remote_region = SqsUrlParser.get_region(queue_url)
+    else:
+        for arn_attribute in ARN_ATTRIBUTES:
+            if is_key_present(span, arn_attribute):
+                arn = span.attributes.get(arn_attribute)
+                remote_account_id = RegionalResourceArnParser.get_account_id(arn)
+                remote_region = RegionalResourceArnParser.get_region(arn)
+                break
+
+    if remote_account_id is not None and remote_region is not None:
+        attributes[AWS_REMOTE_RESOURCE_ACCOUNT_ID] = remote_account_id
+        attributes[AWS_REMOTE_RESOURCE_REGION] = remote_region
+        return True
+    return False
+
+
+def _set_remote_access_key_and_region(span: ReadableSpan, attributes: BoundedAttributes) -> None:
+    if is_key_present(span, AWS_AUTH_ACCESS_KEY):
+        attributes[AWS_REMOTE_RESOURCE_ACCESS_KEY] = span.attributes.get(AWS_AUTH_ACCESS_KEY)
+    if is_key_present(span, AWS_AUTH_REGION):
+        attributes[AWS_REMOTE_RESOURCE_REGION] = span.attributes.get(AWS_AUTH_REGION)
 
 
 def _set_remote_environment(span: ReadableSpan, attributes: BoundedAttributes) -> None:

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

@@ -0,0 +1,45 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Optional
+
+
+class RegionalResourceArnParser:
+    @staticmethod
+    def get_account_id(arn: str) -> Optional[str]:
+        if _is_arn(arn):
+            return str(arn).split(":")[4]
+        return None
+
+    @staticmethod
+    def get_region(arn: str) -> Optional[str]:
+        if _is_arn(arn):
+            return str(arn).split(":")[3]
+        return None
+
+
+def _is_arn(arn: str) -> bool:
+    # Check if arn follows the format:
+    # arn:partition:service:region:account-id:resource-type/resource-id or
+    # arn:partition:service:region:account-id:resource-type:resource-id
+    if arn is None:
+        return False
+
+    if not str(arn).startswith("arn"):
+        return False
+
+    arn_parts = str(arn).split(":")
+    return len(arn_parts) >= 6 and _is_account_id(arn_parts[4])
+
+
+def _is_account_id(input: str) -> bool:
+    if input is None or len(input) != 12:
+        return False
+
+    if not _check_digits(input):
+        return False
+
+    return True
+
+
+def _check_digits(string: str) -> bool:
+    return string.isdigit()

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

@@ -23,6 +23,50 @@ def get_queue_name(url: str) -> Optional[str]:
             return split_url[2]
         return None
 
+    @staticmethod
+    def get_account_id(url: str) -> Optional[str]:
+        """
+        Extracts the account ID from an SQS URL.
+        """
+        if url is None:
+            return None
+        url = url.replace(_HTTP_SCHEMA, "").replace(_HTTPS_SCHEMA, "")
+        split_url: List[Optional[str]] = url.split("/")
+        if _is_valid_sqs_url(url):
+            return split_url[1]
+        return None
+
+    @staticmethod
+    def get_region(url: str) -> Optional[str]:
+        """
+        Extracts the region from an SQS URL.
+        """
+        if url is None:
+            return None
+        url = url.replace(_HTTP_SCHEMA, "").replace(_HTTPS_SCHEMA, "")
+        split_url: List[Optional[str]] = url.split("/")
+        if _is_valid_sqs_url(url):
+            domain: str = split_url[0]
+            domain_parts: List[str] = domain.split(".")
+            if len(domain_parts) == 4:
+                return domain_parts[1]
+        return None
+
+
+def _is_valid_sqs_url(url: str) -> bool:
+    """
+    Checks if the URL is a valid SQS URL.
+    """
+    if url is None:
+        return False
+    split_url: List[str] = url.split("/")
+    return (
+        len(split_url) == 3
+        and split_url[0].lower().startswith("sqs")
+        and _is_account_id(split_url[1])
+        and _is_valid_queue_name(split_url[2])
+    )
+
 
 def _is_account_id(input_str: str) -> bool:
     if input_str is None or len(input_str) != 12: