update url & arn parser; update unit test after version upgrade

Author: blairhyy-amazon

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

@@ -417,16 +417,20 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
             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/", "")
+            remote_resource_identifier = _escape_delimiters(
+                RegionalResourceArnParser.extract_dynamodb_table_name_from_arn(
+                    span.attributes.get(AWS_DYNAMODB_TABLE_ARN)
+                )
             )
         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/", "")
+            remote_resource_identifier = _escape_delimiters(
+                RegionalResourceArnParser.extract_kinesis_stream_name_from_arn(
+                    span.attributes.get(AWS_KINESIS_STREAM_ARN)
+                )
             )
         elif is_key_present(span, _AWS_BUCKET_NAME):
             remote_resource_type = _NORMALIZED_S3_SERVICE_NAME + "::Bucket"
@@ -464,27 +468,35 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
             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]
+            remote_resource_identifier = _escape_delimiters(
+                RegionalResourceArnParser.extract_resource_name_from_arn(
+                    span.attributes.get(AWS_SECRETSMANAGER_SECRET_ARN)
+                )
+            )
             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]
+            remote_resource_identifier = _escape_delimiters(
+                RegionalResourceArnParser.extract_resource_name_from_arn(span.attributes.get(AWS_SNS_TOPIC_ARN))
+            )
             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]
+                RegionalResourceArnParser.extract_resource_name_from_arn(
+                    span.attributes.get(AWS_STEPFUNCTIONS_STATEMACHINE_ARN)
+                )
+            )
             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]
+            remote_resource_identifier = _escape_delimiters(
+                RegionalResourceArnParser.extract_resource_name_from_arn(
+                    span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN)
+                )
+            )
             cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN))
         elif is_key_present(span, AWS_LAMBDA_FUNCTION_NAME):
             # For non-Invoke Lambda operations, treat Lambda as a resource,

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

@@ -75,3 +75,15 @@ def get_aws_region() -> Optional[str]:
     """
     botocore_session = get_aws_session()
     return botocore_session.get_config_variable("region") if botocore_session else None
+
+
+def is_account_id(input_str: str) -> bool:
+    if input_str is None:
+        return False
+
+    try:
+        int(input_str)
+    except ValueError:
+        return False
+
+    return True

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

@@ -2,6 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from typing import Optional
 
+from amazon.opentelemetry.distro._utils import is_account_id
+
 
 class RegionalResourceArnParser:
     @staticmethod
@@ -16,6 +18,28 @@ def get_region(arn: str) -> Optional[str]:
             return str(arn).split(":")[3]
         return None
 
+    @staticmethod
+    def extract_dynamodb_table_name_from_arn(arn: str) -> Optional[str]:
+        resource_name = RegionalResourceArnParser.extract_resource_name_from_arn(arn)
+        if resource_name:
+            return resource_name.replace("table/", "")
+        return None
+
+    @staticmethod
+    def extract_kinesis_stream_name_from_arn(arn: str) -> Optional[str]:
+        resource_name = RegionalResourceArnParser.extract_resource_name_from_arn(arn)
+        if resource_name:
+            return resource_name.replace("stream/", "")
+        return None
+
+    @staticmethod
+    def extract_resource_name_from_arn(arn: str) -> Optional[str]:
+        # Extracts the name of the resource from an arn
+        if _is_arn(arn):
+            split = arn.split(":")
+            return split[-1]
+        return None
+
 
 def _is_arn(arn: str) -> bool:
     # Check if arn follows the format:
@@ -28,18 +52,4 @@ def _is_arn(arn: str) -> bool:
         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()
+    return len(arn_parts) >= 6 and is_account_id(arn_parts[4])

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

@@ -1,6 +1,8 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-from typing import List, Optional
+from typing import List, Optional, Tuple
+
+from amazon.opentelemetry.distro._utils import is_account_id
 
 _HTTP_SCHEMA: str = "http://"
 _HTTPS_SCHEMA: str = "https://"
@@ -12,14 +14,14 @@ def get_queue_name(url: str) -> Optional[str]:
         """
         Best-effort logic to extract queue name from an HTTP url. This method should only be used with a string that is,
         with reasonably high confidence, an SQS queue URL. Handles new/legacy/some custom URLs. Essentially, we require
-        that the URL should have exactly three parts, delimited by /'s (excluding schema), the second part should be a
-        12-digit account id, and the third part should be a valid queue name, per SQS naming conventions.
+        that the URL should have exactly three parts, delimited by /'s (excluding schema), the second part should be an
+        account id consisting of digits, and the third part should be a valid queue name, per SQS naming conventions.
         """
         if url is None:
             return None
-        url = url.replace(_HTTP_SCHEMA, "").replace(_HTTPS_SCHEMA, "")
-        split_url: List[Optional[str]] = url.split("/")
-        if len(split_url) == 3 and _is_account_id(split_url[1]) and _is_valid_queue_name(split_url[2]):
+        urlWithoutProtocol = url.replace(_HTTP_SCHEMA, "").replace(_HTTPS_SCHEMA, "")
+        split_url: List[Optional[str]] = urlWithoutProtocol.split("/")
+        if len(split_url) == 3 and is_account_id(split_url[1]) and _is_valid_queue_name(split_url[2]):
             return split_url[2]
         return None
 
@@ -28,56 +30,38 @@ 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
+        return SqsUrlParser.parse_url(url)[1]
 
     @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])
-    )
+        return SqsUrlParser.parse_url(url)[2]
 
+    @staticmethod
+    def parse_url(url: str) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+        """
+        Parses an SQS URL and extracts its components.
+        URL Format: https://sqs.<region>.amazonaws.com/<accountId>/<queueName>
+        """
+        if url is None:
+            return None, None, None
 
-def _is_account_id(input_str: str) -> bool:
-    if input_str is None or len(input_str) != 12:
-        return False
+        urlWithoutProtocol = url.replace(_HTTP_SCHEMA, "").replace(_HTTPS_SCHEMA, "")
+        split_url: List[Optional[str]] = urlWithoutProtocol.split("/")
+        if (
+            len(split_url) != 3
+            or not is_account_id(split_url[1])
+            or not _is_valid_queue_name(split_url[2])
+            or not split_url[0].lower().startswith("sqs")
+        ):
+            return None, None, None
 
-    try:
-        int(input_str)
-    except ValueError:
-        return False
+        domain: str = split_url[0]
+        domain_parts: List[str] = domain.split(".")
 
-    return True
+        return split_url[2], split_url[1], domain_parts[1] if len(domain_parts) == 4 else None
 
 
 def _is_valid_queue_name(input_str: str) -> bool:

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

@@ -1527,21 +1527,14 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
             keys,
             values,
         )
-        self._validate_remote_resource_attributes(
-            "AWS::StepFunctions::StateMachine",
-            "invalid_arn",
-            None,
-            _AWS_REMOTE_RESOURCE_REGION,
-            None,
-            _AWS_REMOTE_RESOURCE_ACCESS_KEY,
-        )
+        self._validate_remote_resource_attributes(None, None, None)
         self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None])
 
         # Invalid arn and no account access key
         self._mock_attribute(
             [AWS_STEPFUNCTIONS_STATEMACHINE_ARN, SpanAttributes.RPC_SYSTEM], ["invalid_arn", "aws-api"]
         )
-        self._validate_remote_resource_attributes("AWS::StepFunctions::StateMachine", "invalid_arn")
+        self._validate_remote_resource_attributes(None, None, None)
         self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None])
 
         # Both account access key and account id are not available
@@ -1571,11 +1564,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
             [AWS_STEPFUNCTIONS_STATEMACHINE_ARN, SpanAttributes.RPC_SYSTEM],
             ["arn:aws:states:us-east-1:invalid_account_id:stateMachine:testStateMachine", "aws-api"],
         )
-        self._validate_remote_resource_attributes(
-            "AWS::StepFunctions::StateMachine",
-            "testStateMachine",
-            "arn:aws:states:us-east-1:invalid_account_id:stateMachine:testStateMachine",
-        )
+        self._validate_remote_resource_attributes(None, None, None)
         self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None])
 
         # Arn with invalid region