diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py index a2b7cd00b..5568b4d6a 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py +++ b/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. diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py index 235ef47e0..efa8661c1 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py +++ b/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: diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py index b040c07ab..edc1b19ec 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py +++ b/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 = { diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py index 870ad6599..a9ff24698 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py +++ b/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) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py index be17c4402..16004c726 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py +++ b/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])