Skip to content

Commit 41e0987

Browse files
authored
feat: Add CFN Primary ID for Existing Supported AWS Resources (#264)
### *Description of changes:* These changes support AWS Resources with different CFN Identifier formats. This new attribute will have the key `aws.remote.resource.cfn.primary.identifier`. More context can be found in the comments of this closed PR: #261 In most cases, this new attribute should have the same value as the existing attribute `aws.remote.resource.identifier`. However, there are some edge cases we must handle as seen in this PR: - `AWS::SQS::Queue` - https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-sqs-queue.html - The CFN Id should be the URL of the Queue. - `AWS::Bedrock::DataSource` - https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-bedrock-datasource.html - The CFN Id should be a compound value where the knowledge base Id and the data source Id are separated by `|` Also updated Bedrock Agent request parameters to also include Knowledge Base ID: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent/client/get_data_source.html ### *Test Plan:* Set up a client-server with auto-instrumentation to verify that the correct span data is being generated. ![bedrock_cfn_primary_id_verification](https://github.com/user-attachments/assets/9a375751-a70c-41ab-896a-00e26a5fcc76) ![sqs_cfn_primary_id_verification](https://github.com/user-attachments/assets/a8ee28b0-4344-4f48-82b7-a100da939860) Unit tests for instrumentation ![metric_generator_unit_test_verification](https://github.com/user-attachments/assets/52c26651-068d-4d40-a020-69714feab17c) ![instrumentation_patch_unit_test_verification](https://github.com/user-attachments/assets/3c23031d-52c2-49e6-9285-da679351ae0d) By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.
1 parent b1fcdb3 commit 41e0987

File tree

5 files changed

+49
-7
lines changed

5 files changed

+49
-7
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
AWS_SDK_DESCENDANT: str = "aws.sdk.descendant"
1313
AWS_CONSUMER_PARENT_SPAN_KIND: str = "aws.consumer.parent.span.kind"
1414
AWS_TRACE_FLAG_SAMPLED: str = "aws.trace.flag.sampled"
15+
AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER: str = "aws.remote.resource.cfn.primary.identifier"
1516

1617
# AWS_#_NAME attributes are not supported in python as they are not part of the Semantic Conventions.
1718
# TODO:Move to Semantic Conventions when these attributes are added.

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

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
AWS_BEDROCK_DATA_SOURCE_ID,
1111
AWS_BEDROCK_GUARDRAIL_ID,
1212
AWS_BEDROCK_KNOWLEDGE_BASE_ID,
13+
AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
1314
AWS_KINESIS_STREAM_NAME,
1415
AWS_LOCAL_OPERATION,
1516
AWS_LOCAL_SERVICE,
@@ -372,6 +373,7 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
372373
"""
373374
remote_resource_type: Optional[str] = None
374375
remote_resource_identifier: Optional[str] = None
376+
cloudformation_primary_identifier: Optional[str] = None
375377

376378
if is_aws_sdk_span(span):
377379
# 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
387389
elif is_key_present(span, AWS_SQS_QUEUE_NAME):
388390
remote_resource_type = _NORMALIZED_SQS_SERVICE_NAME + "::Queue"
389391
remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_SQS_QUEUE_NAME))
392+
cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_SQS_QUEUE_URL))
390393
elif is_key_present(span, AWS_SQS_QUEUE_URL):
391394
remote_resource_type = _NORMALIZED_SQS_SERVICE_NAME + "::Queue"
392395
remote_resource_identifier = _escape_delimiters(
393396
SqsUrlParser.get_queue_name(span.attributes.get(AWS_SQS_QUEUE_URL))
394397
)
398+
cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_SQS_QUEUE_URL))
395399
elif is_key_present(span, AWS_BEDROCK_AGENT_ID):
396400
remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Agent"
397401
remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_AGENT_ID))
398402
elif is_key_present(span, AWS_BEDROCK_DATA_SOURCE_ID):
399403
remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::DataSource"
400404
remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_DATA_SOURCE_ID))
405+
cloudformation_primary_identifier = (
406+
_escape_delimiters(span.attributes.get(AWS_BEDROCK_KNOWLEDGE_BASE_ID))
407+
+ "|"
408+
+ remote_resource_identifier
409+
)
401410
elif is_key_present(span, AWS_BEDROCK_GUARDRAIL_ID):
402411
remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Guardrail"
403412
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
411420
remote_resource_type = _DB_CONNECTION_STRING_TYPE
412421
remote_resource_identifier = _get_db_connection(span)
413422

414-
if remote_resource_type is not None and remote_resource_identifier is not None:
423+
# If the CFN Primary Id is still None here, that means it is not an edge case.
424+
# Then, we can just assign it the same value as remote_resource_identifier
425+
if cloudformation_primary_identifier is None:
426+
cloudformation_primary_identifier = remote_resource_identifier
427+
428+
if (
429+
remote_resource_type is not None
430+
and remote_resource_identifier is not None
431+
and cloudformation_primary_identifier is not None
432+
):
415433
attributes[AWS_REMOTE_RESOURCE_TYPE] = remote_resource_type
416434
attributes[AWS_REMOTE_RESOURCE_IDENTIFIER] = remote_resource_identifier
435+
attributes[AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER] = cloudformation_primary_identifier
417436

418437

419438
def _get_db_connection(span: ReadableSpan) -> None:

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ class _DataSourceOperation(_BedrockAgentOperation):
125125
"""
126126

127127
request_attributes = {
128+
AWS_BEDROCK_KNOWLEDGE_BASE_ID: _KNOWLEDGE_BASE_ID,
128129
AWS_BEDROCK_DATA_SOURCE_ID: _DATA_SOURCE_ID,
129130
}
130131
response_attributes = {

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

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1044,14 +1044,24 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
10441044
self._mock_attribute([AWS_BEDROCK_AGENT_ID], [None])
10451045

10461046
# Validate behaviour of AWS_BEDROCK_DATA_SOURCE_ID attribute, then remove it.
1047-
self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID], ["test_datasource_id"], keys, values)
1047+
self._mock_attribute(
1048+
[AWS_BEDROCK_DATA_SOURCE_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID],
1049+
["test_datasource_id", "test_knowledge_base_id"],
1050+
keys,
1051+
values,
1052+
)
10481053
self._validate_remote_resource_attributes("AWS::Bedrock::DataSource", "test_datasource_id")
1049-
self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID], [None])
1054+
self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID], [None, None])
10501055

10511056
# Validate behaviour of AWS_BEDROCK_DATA_SOURCE_ID attribute with special chars(^), then remove it.
1052-
self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID], ["test_datasource_^id"], keys, values)
1057+
self._mock_attribute(
1058+
[AWS_BEDROCK_DATA_SOURCE_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID],
1059+
["test_datasource_^id", "test_knowledge_base_^id"],
1060+
keys,
1061+
values,
1062+
)
10531063
self._validate_remote_resource_attributes("AWS::Bedrock::DataSource", "test_datasource_^^id")
1054-
self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID], [None])
1064+
self._mock_attribute([AWS_BEDROCK_DATA_SOURCE_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID], [None, None])
10551065

10561066
# Validate behaviour of AWS_BEDROCK_GUARDRAIL_ID attribute, then remove it.
10571067
self._mock_attribute([AWS_BEDROCK_GUARDRAIL_ID], ["test_guardrail_id"], keys, values)

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

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -292,12 +292,23 @@ def _test_patched_bedrock_agent_instrumentation(self):
292292
"UpdateDataSource": ("aws.bedrock.data_source.id", _BEDROCK_DATASOURCE_ID),
293293
}
294294

295+
data_source_operations = ["DeleteDataSource", "GetDataSource", "UpdateDataSource"]
296+
295297
for operation, attribute_tuple in operation_to_expected_attribute.items():
296298
bedrock_agent_extract_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
297299
"bedrock-agent", operation
298300
)
299-
self.assertEqual(len(bedrock_agent_extract_attributes), 1)
300-
self.assertEqual(bedrock_agent_extract_attributes[attribute_tuple[0]], attribute_tuple[1])
301+
302+
if operation in data_source_operations:
303+
self.assertEqual(len(bedrock_agent_extract_attributes), 2)
304+
self.assertEqual(bedrock_agent_extract_attributes[attribute_tuple[0]], attribute_tuple[1])
305+
self.assertEqual(
306+
bedrock_agent_extract_attributes["aws.bedrock.knowledge_base.id"], _BEDROCK_KNOWLEDGEBASE_ID
307+
)
308+
else:
309+
self.assertEqual(len(bedrock_agent_extract_attributes), 1)
310+
self.assertEqual(bedrock_agent_extract_attributes[attribute_tuple[0]], attribute_tuple[1])
311+
301312
bedrock_agent_success_attributes: Dict[str, str] = _do_on_success_bedrock("bedrock-agent", operation)
302313
self.assertEqual(len(bedrock_agent_success_attributes), 1)
303314
self.assertEqual(bedrock_agent_success_attributes[attribute_tuple[0]], attribute_tuple[1])

0 commit comments

Comments
 (0)