feat: Add CFN Primary Id for Existing Supported AWS Resources

Author: yiyuan-he

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

@@ -12,6 +12,7 @@
 AWS_SDK_DESCENDANT: str = "aws.sdk.descendant"
 AWS_CONSUMER_PARENT_SPAN_KIND: str = "aws.consumer.parent.span.kind"
 AWS_TRACE_FLAG_SAMPLED: str = "aws.trace.flag.sampled"
+AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER: str = "aws.remote.cfn.primary.identifier"
 
 # AWS_#_NAME attributes are not supported in python as they are not part of the Semantic Conventions.
 # TODO：Move to Semantic Conventions when these attributes are added.

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

@@ -10,6 +10,7 @@
     AWS_BEDROCK_DATA_SOURCE_ID,
     AWS_BEDROCK_GUARDRAIL_ID,
     AWS_BEDROCK_KNOWLEDGE_BASE_ID,
+    AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
     AWS_KINESIS_STREAM_NAME,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
@@ -372,6 +373,7 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
     """
     remote_resource_type: Optional[str] = None
     remote_resource_identifier: Optional[str] = None
+    cloudformation_primary_identifier: Optional[str] = None
 
     if is_aws_sdk_span(span):
         # Only extract the table name when _AWS_TABLE_NAMES has size equals to one
@@ -387,17 +389,24 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
         elif is_key_present(span, AWS_SQS_QUEUE_NAME):
             remote_resource_type = _NORMALIZED_SQS_SERVICE_NAME + "::Queue"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_SQS_QUEUE_NAME))
+            cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_SQS_QUEUE_URL))
         elif is_key_present(span, AWS_SQS_QUEUE_URL):
             remote_resource_type = _NORMALIZED_SQS_SERVICE_NAME + "::Queue"
             remote_resource_identifier = _escape_delimiters(
                 SqsUrlParser.get_queue_name(span.attributes.get(AWS_SQS_QUEUE_URL))
             )
+            cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_SQS_QUEUE_URL))
         elif is_key_present(span, AWS_BEDROCK_AGENT_ID):
             remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Agent"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_AGENT_ID))
         elif is_key_present(span, AWS_BEDROCK_DATA_SOURCE_ID):
             remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::DataSource"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_DATA_SOURCE_ID))
+            cloudformation_primary_identifier = (
+                _escape_delimiters(span.attributes.get(AWS_BEDROCK_KNOWLEDGE_BASE_ID))
+                + "|"
+                + remote_resource_identifier
+            )
         elif is_key_present(span, AWS_BEDROCK_GUARDRAIL_ID):
             remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Guardrail"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_GUARDRAIL_ID))
@@ -411,9 +420,19 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
         remote_resource_type = _DB_CONNECTION_STRING_TYPE
         remote_resource_identifier = _get_db_connection(span)
 
-    if remote_resource_type is not None and remote_resource_identifier is not None:
+    # If the CFN Primary Id is still None here, that means it is not an edge case.
+    # Then, we can just assign it the same value as remote_resource_identifier
+    if cloudformation_primary_identifier is None:
+        cloudformation_primary_identifier = remote_resource_identifier
+
+    if (
+        remote_resource_type is not None
+        and remote_resource_identifier is not None
+        and cloudformation_primary_identifier is not None
+    ):
         attributes[AWS_REMOTE_RESOURCE_TYPE] = remote_resource_type
         attributes[AWS_REMOTE_RESOURCE_IDENTIFIER] = remote_resource_identifier
+        attributes[AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER] = cloudformation_primary_identifier
 
 
 def _get_db_connection(span: ReadableSpan) -> None:

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

@@ -125,6 +125,7 @@ class _DataSourceOperation(_BedrockAgentOperation):
     """
 
     request_attributes = {
+        AWS_BEDROCK_KNOWLEDGE_BASE_ID: _KNOWLEDGE_BASE_ID,
         AWS_BEDROCK_DATA_SOURCE_ID: _DATA_SOURCE_ID,
     }
     response_attributes = {

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

@@ -1044,14 +1044,24 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         self._mock_attribute([AWS_BEDROCK_AGENT_ID], [None])
 
         # Validate behaviour of AWS_BEDROCK_DATA_SOURCE_ID attribute, then remove it.
-        self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID], ["test_datasource_id"], keys, values)
+        self._mock_attribute(
+            [AWS_BEDROCK_DATA_SOURCE_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID],
+            ["test_datasource_id", "test_knowledge_base_id"],
+            keys,
+            values,
+        )
         self._validate_remote_resource_attributes("AWS::Bedrock::DataSource", "test_datasource_id")
-        self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID], [None])
+        self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID], [None, None])
 
         # Validate behaviour of AWS_BEDROCK_DATA_SOURCE_ID attribute with special chars(^), then remove it.
-        self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID], ["test_datasource_^id"], keys, values)
+        self._mock_attribute(
+            [AWS_BEDROCK_DATA_SOURCE_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID],
+            ["test_datasource_^id", "test_knowledge_base_^id"],
+            keys,
+            values,
+        )
         self._validate_remote_resource_attributes("AWS::Bedrock::DataSource", "test_datasource_^^id")
-        self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID], [None])
+        self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID], [None, None])
 
         # Validate behaviour of AWS_BEDROCK_GUARDRAIL_ID attribute, then remove it.
         self._mock_attribute([AWS_BEDROCK_GUARDRAIL_ID], ["test_guardrail_id"], keys, values)

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

@@ -292,12 +292,23 @@ def _test_patched_bedrock_agent_instrumentation(self):
             "UpdateDataSource": ("aws.bedrock.data_source.id", _BEDROCK_DATASOURCE_ID),
         }
 
+        data_source_operations = ["DeleteDataSource", "GetDataSource", "UpdateDataSource"]
+
         for operation, attribute_tuple in operation_to_expected_attribute.items():
             bedrock_agent_extract_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
                 "bedrock-agent", operation
             )
-            self.assertEqual(len(bedrock_agent_extract_attributes), 1)
-            self.assertEqual(bedrock_agent_extract_attributes[attribute_tuple[0]], attribute_tuple[1])
+
+            if operation in data_source_operations:
+                self.assertEqual(len(bedrock_agent_extract_attributes), 2)
+                self.assertEqual(bedrock_agent_extract_attributes[attribute_tuple[0]], attribute_tuple[1])
+                self.assertEqual(
+                    bedrock_agent_extract_attributes["aws.bedrock.knowledge_base.id"], _BEDROCK_KNOWLEDGEBASE_ID
+                )
+            else:
+                self.assertEqual(len(bedrock_agent_extract_attributes), 1)
+                self.assertEqual(bedrock_agent_extract_attributes[attribute_tuple[0]], attribute_tuple[1])
+
             bedrock_agent_success_attributes: Dict[str, str] = _do_on_success_bedrock("bedrock-agent", operation)
             self.assertEqual(len(bedrock_agent_success_attributes), 1)
             self.assertEqual(bedrock_agent_success_attributes[attribute_tuple[0]], attribute_tuple[1])