From 3dd31c5163f2a70abe6976b6dc9e49ec88f93a4e Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Wed, 25 Sep 2024 11:45:29 -0700 Subject: [PATCH 01/11] feat: Add auto-instrumentation support for SecretsManager --- .../distro/_aws_attribute_keys.py | 1 + .../distro/_aws_metric_attribute_generator.py | 6 ++++ .../distro/patches/_botocore_patches.py | 34 ++++++++++++++++++- .../test_aws_metric_attribute_generator.py | 14 ++++++++ .../distro/test_instrumentation_patch.py | 25 ++++++++++++++ 5 files changed, 79 insertions(+), 1 deletion(-) 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 6ecd8a846..16fb97a0e 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 @@ -23,3 +23,4 @@ AWS_BEDROCK_KNOWLEDGE_BASE_ID: str = "aws.bedrock.knowledge_base.id" AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id" AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail.id" +AWS_SECRETSMANAGER_SECRET_ARN: str = "aws.secretsmanager.secret.arn" 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 efa8661c1..4e35660be 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 @@ -19,6 +19,7 @@ AWS_REMOTE_RESOURCE_IDENTIFIER, AWS_REMOTE_RESOURCE_TYPE, AWS_REMOTE_SERVICE, + AWS_SECRETSMANAGER_SECRET_ARN, AWS_SPAN_KIND, AWS_SQS_QUEUE_NAME, AWS_SQS_QUEUE_URL, @@ -88,6 +89,7 @@ _NORMALIZED_SQS_SERVICE_NAME: str = "AWS::SQS" _NORMALIZED_BEDROCK_SERVICE_NAME: str = "AWS::Bedrock" _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME: str = "AWS::BedrockRuntime" +_NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager" _DB_CONNECTION_STRING_TYPE: str = "DB::Connection" # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present. @@ -309,6 +311,7 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str "Bedrock Agent": _NORMALIZED_BEDROCK_SERVICE_NAME, "Bedrock Agent Runtime": _NORMALIZED_BEDROCK_SERVICE_NAME, "Bedrock Runtime": _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME, + "Secrets Manager": _NORMALIZED_SECRETSMANAGER_SERVICE_NAME, } return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name) return service_name @@ -416,6 +419,9 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri elif is_key_present(span, GEN_AI_REQUEST_MODEL): remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Model" 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)) elif is_db_span(span): remote_resource_type = _DB_CONNECTION_STRING_TYPE remote_resource_identifier = _get_db_connection(span) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index d635f7d1b..9e07806da 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -5,6 +5,7 @@ from amazon.opentelemetry.distro._aws_attribute_keys import ( AWS_KINESIS_STREAM_NAME, + AWS_SECRETSMANAGER_SECRET_ARN, AWS_SQS_QUEUE_NAME, AWS_SQS_QUEUE_URL, ) @@ -16,8 +17,9 @@ ) from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS from opentelemetry.instrumentation.botocore.extensions.sqs import _SqsExtension -from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension +from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension, _BotoResultT from opentelemetry.semconv.trace import SpanAttributes +from opentelemetry.trace.span import Span def _apply_botocore_instrumentation_patches() -> None: @@ -29,6 +31,19 @@ def _apply_botocore_instrumentation_patches() -> None: _apply_botocore_s3_patch() _apply_botocore_sqs_patch() _apply_botocore_bedrock_patch() + _apply_botocore_secretsmanager_patch() + + +def _apply_botocore_secretsmanager_patch() -> None: + """Botocore instrumentation patch for SecretsManager + + This patch adds an extension to the upstream's list of known extension for SecretsManager. + Extensions allow for custom logic for adding service-specific information to spans, such as + attributes. Specifically, we are adding logic to add the `aws.secretsmanager.secret.arn` + attribute, to be used to generate RemoteTarget and achieve parity with the Java + instrumentation. + """ + _KNOWN_EXTENSIONS["secretsmanager"] = _lazy_load(".", "_SecretsManagerExtension") def _apply_botocore_kinesis_patch() -> None: @@ -108,6 +123,23 @@ def loader(): # END The OpenTelemetry Authors code +class _SecretsManagerExtension(_AwsSdkExtension): + def extract_attributes(self, attributes: _AttributeMapT): + """ + SecretId can be secret name or secret arn, the function extracts attributes + only if the SecretId parameter is provided as an arn which starts with + `arn:aws:secretsmanager:` + """ + secret_id = self._call_context.params.get("SecretId") + if secret_id and secret_id.startswith("arn:aws:secretsmanager:"): + attributes[AWS_SECRETSMANAGER_SECRET_ARN] = secret_id + + def on_success(self, span: Span, result: _BotoResultT): + secret_arn = result.get("ARN") + if secret_arn: + span.set_attribute(AWS_SECRETSMANAGER_SECRET_ARN, secret_arn) + + class _S3Extension(_AwsSdkExtension): def extract_attributes(self, attributes: _AttributeMapT): bucket_name = self._call_context.params.get("Bucket") 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 a9ff24698..22e6c6d66 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 @@ -21,6 +21,7 @@ AWS_REMOTE_RESOURCE_IDENTIFIER, AWS_REMOTE_RESOURCE_TYPE, AWS_REMOTE_SERVICE, + AWS_SECRETSMANAGER_SECRET_ARN, AWS_SPAN_KIND, AWS_SQS_QUEUE_NAME, AWS_SQS_QUEUE_URL, @@ -877,6 +878,7 @@ def test_normalize_remote_service_name_aws_sdk(self): self.validate_aws_sdk_service_normalization("Bedrock Agent", "AWS::Bedrock") self.validate_aws_sdk_service_normalization("Bedrock Agent Runtime", "AWS::Bedrock") self.validate_aws_sdk_service_normalization("Bedrock Runtime", "AWS::BedrockRuntime") + self.validate_aws_sdk_service_normalization("Secrets Manager", "AWS::SecretsManager") def validate_aws_sdk_service_normalization(self, service_name: str, expected_remote_service: str): self._mock_attribute([SpanAttributes.RPC_SYSTEM, SpanAttributes.RPC_SERVICE], ["aws-api", service_name]) @@ -1093,6 +1095,18 @@ def test_sdk_client_span_with_remote_resource_attributes(self): self._validate_remote_resource_attributes("AWS::Bedrock::Model", "test.service_^^id") self._mock_attribute([GEN_AI_REQUEST_MODEL], [None]) + # Validate behaviour of AWS_SECRETSMANAGER_SECRET_ARN attribute, then remove it. + self._mock_attribute( + [AWS_SECRETSMANAGER_SECRET_ARN], + ["arn:aws:secretsmanager:us-east-1:123456789012:secret:secret_name-lERW9H"], + keys, + values + ) + self._validate_remote_resource_attributes( + "AWS::SecretsManager::Secret", "arn:aws:secretsmanager:us-east-1:123456789012:secret:secret_name-lERW9H" + ) + self._mock_attribute([AWS_SECRETSMANAGER_SECRET_ARN], [None]) + self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None]) def test_client_db_span_with_remote_resource_attributes(self): 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 16004c726..65af6be8f 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 @@ -26,6 +26,7 @@ _BEDROCK_KNOWLEDGEBASE_ID: str = "KnowledgeBaseId" _GEN_AI_SYSTEM: str = "aws_bedrock" _GEN_AI_REQUEST_MODEL: str = "genAiReuqestModelId" +_SECRET_ARN: str = "arn:aws:secretsmanager:us-west-2:000000000000:secret:testSecret-ABCDEF" # Patch names GET_DISTRIBUTION_PATCH: str = ( @@ -141,6 +142,9 @@ def _test_unpatched_botocore_instrumentation(self): # BedrockRuntime self.assertFalse("bedrock-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a bedrock-runtime extension") + # SecretsManager + self.assertFalse("secretsmanager" in _KNOWN_EXTENSIONS, "Upstream has added a SecretsManager extension") + def _test_unpatched_gevent_instrumentation(self): self.assertFalse(gevent.monkey.is_module_patched("os"), "gevent os module has been patched") self.assertFalse(gevent.monkey.is_module_patched("thread"), "gevent thread module has been patched") @@ -200,6 +204,15 @@ def _test_patched_botocore_instrumentation(self): self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM) self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], _GEN_AI_REQUEST_MODEL) + # SecretsManager + self.assertTrue("secretsmanager" in _KNOWN_EXTENSIONS) + secretsmanager_attributes: Dict[str, str] = _do_extract_secretsmanager_attributes() + self.assertTrue("aws.secretsmanager.secret.arn" in secretsmanager_attributes) + self.assertEqual(secretsmanager_attributes["aws.secretsmanager.secret.arn"], _SECRET_ARN) + secretsmanager_success_attributes: Dict[str, str] = _do_on_success_secretsmanager() + self.assertTrue("aws.secretsmanager.secret.arn" in secretsmanager_success_attributes) + self.assertEqual(secretsmanager_success_attributes["aws.secretsmanager.secret.arn"], _SECRET_ARN) + def _test_patched_gevent_os_ssl_instrumentation(self): # Only ssl and os module should have been patched since the environment variable was set to 'os, ssl' self.assertTrue(gevent.monkey.is_module_patched("ssl"), "gevent ssl module has not been patched") @@ -358,6 +371,18 @@ def _do_on_success_bedrock(service, operation=None) -> Dict[str, str]: return _do_on_success(service, result, operation) +def _do_extract_secretsmanager_attributes() -> Dict[str, str]: + service_name: str = "secretsmanager" + params: Dict[str, str] = {"SecretId": _SECRET_ARN} + return _do_extract_attributes(service_name, params) + + +def _do_on_success_secretsmanager() -> Dict[str, str]: + service_name: str = "secretsmanager" + result: Dict[str, Any] = {"ARN": _SECRET_ARN} + return _do_on_success(service_name, result) + + def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]: mock_call_context: MagicMock = MagicMock() mock_call_context.params = params From 6fcf4cc2e04b3f838def4b8db1597f518ff6c31a Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Wed, 25 Sep 2024 12:22:00 -0700 Subject: [PATCH 02/11] feat: Add auto-instrumentation support for SNS --- .../distro/_aws_attribute_keys.py | 1 + .../distro/_aws_metric_attribute_generator.py | 6 +++++ .../distro/patches/_botocore_patches.py | 26 +++++++++++++++++++ .../test_aws_metric_attribute_generator.py | 7 +++++ .../distro/test_instrumentation_patch.py | 16 ++++++++++++ 5 files changed, 56 insertions(+) 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 16fb97a0e..26fa0bee7 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 @@ -24,3 +24,4 @@ AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id" AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail.id" AWS_SECRETSMANAGER_SECRET_ARN: str = "aws.secretsmanager.secret.arn" +AWS_SNS_TOPIC_ARN: str = "aws.sns.topic.arn" 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 4e35660be..e290c3295 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 @@ -20,6 +20,7 @@ AWS_REMOTE_RESOURCE_TYPE, AWS_REMOTE_SERVICE, AWS_SECRETSMANAGER_SECRET_ARN, + AWS_SNS_TOPIC_ARN, AWS_SPAN_KIND, AWS_SQS_QUEUE_NAME, AWS_SQS_QUEUE_URL, @@ -90,6 +91,7 @@ _NORMALIZED_BEDROCK_SERVICE_NAME: str = "AWS::Bedrock" _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME: str = "AWS::BedrockRuntime" _NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager" +_NORMALIZED_SNS_SERVICE_NAME: str = "AWS::SNS" _DB_CONNECTION_STRING_TYPE: str = "DB::Connection" # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present. @@ -312,6 +314,7 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str "Bedrock Agent Runtime": _NORMALIZED_BEDROCK_SERVICE_NAME, "Bedrock Runtime": _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME, "Secrets Manager": _NORMALIZED_SECRETSMANAGER_SERVICE_NAME, + "SNS": _NORMALIZED_SNS_SERVICE_NAME, } return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name) return service_name @@ -422,6 +425,9 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri 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)) + 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)) elif is_db_span(span): remote_resource_type = _DB_CONNECTION_STRING_TYPE remote_resource_identifier = _get_db_connection(span) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index 9e07806da..46b126291 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -6,6 +6,7 @@ from amazon.opentelemetry.distro._aws_attribute_keys import ( AWS_KINESIS_STREAM_NAME, AWS_SECRETSMANAGER_SECRET_ARN, + AWS_SNS_TOPIC_ARN, AWS_SQS_QUEUE_NAME, AWS_SQS_QUEUE_URL, ) @@ -16,6 +17,7 @@ _BedrockRuntimeExtension, ) from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS +from opentelemetry.instrumentation.botocore.extensions.sns import _SnsExtension from opentelemetry.instrumentation.botocore.extensions.sqs import _SqsExtension from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension, _BotoResultT from opentelemetry.semconv.trace import SpanAttributes @@ -32,8 +34,32 @@ def _apply_botocore_instrumentation_patches() -> None: _apply_botocore_sqs_patch() _apply_botocore_bedrock_patch() _apply_botocore_secretsmanager_patch() + _apply_botocore_sns_patch() +def _apply_botocore_sns_patch() -> None: + """Botocore instrumentation patch for SNS + + This patch adds an extension to the upstream's list of known extensions for SNS. + Extensions allow for custom logic for adding service-specific information to + spans, such as attributes. Specifically, we are adding logic to add the + `aws.sns.topic.arn` attribute, to be used to generate RemoteTarget and achieve + parity with the Java instrumentation. + + Sidenote: There exists SpanAttributes.MESSAGING_DESTINATION_NAME in the upstream + logic that we could re-purpose here. We do not use it here to maintain consistent + naming patterns with other AWS resources. + """ + old_extract_attributes = _SnsExtension.extract_attributes + + def patch_extract_attributes(self, attributes: _AttributeMapT): + old_extract_attributes(self, attributes) + topic_arn = self._call_context.params.get("TopicArn") + if topic_arn: + attributes[AWS_SNS_TOPIC_ARN] = topic_arn + + _SnsExtension.extract_attributes = patch_extract_attributes + def _apply_botocore_secretsmanager_patch() -> None: """Botocore instrumentation patch for SecretsManager 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 22e6c6d66..3dcec10db 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 @@ -22,6 +22,7 @@ AWS_REMOTE_RESOURCE_TYPE, AWS_REMOTE_SERVICE, AWS_SECRETSMANAGER_SECRET_ARN, + AWS_SNS_TOPIC_ARN, AWS_SPAN_KIND, AWS_SQS_QUEUE_NAME, AWS_SQS_QUEUE_URL, @@ -879,6 +880,7 @@ def test_normalize_remote_service_name_aws_sdk(self): self.validate_aws_sdk_service_normalization("Bedrock Agent Runtime", "AWS::Bedrock") self.validate_aws_sdk_service_normalization("Bedrock Runtime", "AWS::BedrockRuntime") self.validate_aws_sdk_service_normalization("Secrets Manager", "AWS::SecretsManager") + self.validate_aws_sdk_service_normalization("SNS", "AWS::SNS") def validate_aws_sdk_service_normalization(self, service_name: str, expected_remote_service: str): self._mock_attribute([SpanAttributes.RPC_SYSTEM, SpanAttributes.RPC_SERVICE], ["aws-api", service_name]) @@ -1107,6 +1109,11 @@ def test_sdk_client_span_with_remote_resource_attributes(self): ) self._mock_attribute([AWS_SECRETSMANAGER_SECRET_ARN], [None]) + # Validate behaviour of AWS_SNS_TOPIC_ARN attribute, then remove it. + self._mock_attribute([AWS_SNS_TOPIC_ARN], ["arn:aws:sns:us-west-2:012345678901:test_topic"], keys, values) + self._validate_remote_resource_attributes("AWS::SNS::Topic", "arn:aws:sns:us-west-2:012345678901:test_topic") + self._mock_attribute([AWS_SNS_TOPIC_ARN], [None]) + self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None]) def test_client_db_span_with_remote_resource_attributes(self): 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 65af6be8f..25b191670 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 @@ -27,6 +27,7 @@ _GEN_AI_SYSTEM: str = "aws_bedrock" _GEN_AI_REQUEST_MODEL: str = "genAiReuqestModelId" _SECRET_ARN: str = "arn:aws:secretsmanager:us-west-2:000000000000:secret:testSecret-ABCDEF" +_TOPIC_ARN: str = "topicArn" # Patch names GET_DISTRIBUTION_PATCH: str = ( @@ -145,6 +146,9 @@ def _test_unpatched_botocore_instrumentation(self): # SecretsManager self.assertFalse("secretsmanager" in _KNOWN_EXTENSIONS, "Upstream has added a SecretsManager extension") + # SNS + self.assertTrue("sns" in _KNOWN_EXTENSIONS, "Upstream has removed the SNS extension") + def _test_unpatched_gevent_instrumentation(self): self.assertFalse(gevent.monkey.is_module_patched("os"), "gevent os module has been patched") self.assertFalse(gevent.monkey.is_module_patched("thread"), "gevent thread module has been patched") @@ -213,6 +217,12 @@ def _test_patched_botocore_instrumentation(self): self.assertTrue("aws.secretsmanager.secret.arn" in secretsmanager_success_attributes) self.assertEqual(secretsmanager_success_attributes["aws.secretsmanager.secret.arn"], _SECRET_ARN) + # SNS + self.assertTrue("sns" in _KNOWN_EXTENSIONS) + sns_attributes: Dict[str, str] = _do_extract_sns_attributes() + self.assertTrue("aws.sns.topic.arn" in sns_attributes) + self.assertEqual(sns_attributes["aws.sns.topic.arn"], _TOPIC_ARN) + def _test_patched_gevent_os_ssl_instrumentation(self): # Only ssl and os module should have been patched since the environment variable was set to 'os, ssl' self.assertTrue(gevent.monkey.is_module_patched("ssl"), "gevent ssl module has not been patched") @@ -383,6 +393,12 @@ def _do_on_success_secretsmanager() -> Dict[str, str]: return _do_on_success(service_name, result) +def _do_extract_sns_attributes() -> Dict[str, str]: + service_name: str = "sns" + params: Dict[str, str] = {"TopicArn": _TOPIC_ARN} + return _do_extract_attributes(service_name, params) + + def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]: mock_call_context: MagicMock = MagicMock() mock_call_context.params = params From a8bd2389ce1c7b25167a1003baf9ead0f04a74a9 Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Wed, 25 Sep 2024 13:16:30 -0700 Subject: [PATCH 03/11] feat: Add auto-instrumentation support for StepFunctions --- .../distro/_aws_attribute_keys.py | 2 ++ .../distro/_aws_metric_attribute_generator.py | 10 +++++++ .../distro/patches/_botocore_patches.py | 26 +++++++++++++++++ .../test_aws_metric_attribute_generator.py | 28 +++++++++++++++++++ .../distro/test_instrumentation_patch.py | 19 +++++++++++++ 5 files changed, 85 insertions(+) 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 26fa0bee7..2d865ea1e 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 @@ -25,3 +25,5 @@ AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail.id" AWS_SECRETSMANAGER_SECRET_ARN: str = "aws.secretsmanager.secret.arn" AWS_SNS_TOPIC_ARN: str = "aws.sns.topic.arn" +AWS_STEPFUNCTIONS_STATEMACHINE_ARN: str = "aws.stepfunctions.state_machine.arn" +AWS_STEPFUNCTIONS_ACTIVITY_ARN: str = "aws.stepfunctions.activity.arn" 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 e290c3295..28a6cc116 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 @@ -24,6 +24,8 @@ AWS_SPAN_KIND, AWS_SQS_QUEUE_NAME, AWS_SQS_QUEUE_URL, + AWS_STEPFUNCTIONS_ACTIVITY_ARN, + AWS_STEPFUNCTIONS_STATEMACHINE_ARN, ) from amazon.opentelemetry.distro._aws_span_processing_util import ( GEN_AI_REQUEST_MODEL, @@ -92,6 +94,7 @@ _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME: str = "AWS::BedrockRuntime" _NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager" _NORMALIZED_SNS_SERVICE_NAME: str = "AWS::SNS" +_NORMALIZED_STEPFUNCTIONS_SERVICE_NAME: str = "AWS::StepFunctions" _DB_CONNECTION_STRING_TYPE: str = "DB::Connection" # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present. @@ -315,6 +318,7 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str "Bedrock Runtime": _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME, "Secrets Manager": _NORMALIZED_SECRETSMANAGER_SERVICE_NAME, "SNS": _NORMALIZED_SNS_SERVICE_NAME, + "SFN": _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME, } return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name) return service_name @@ -428,6 +432,12 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri 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)) + 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)) + 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)) elif is_db_span(span): remote_resource_type = _DB_CONNECTION_STRING_TYPE remote_resource_identifier = _get_db_connection(span) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index 46b126291..5a62d6002 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -9,6 +9,8 @@ AWS_SNS_TOPIC_ARN, AWS_SQS_QUEUE_NAME, AWS_SQS_QUEUE_URL, + AWS_STEPFUNCTIONS_ACTIVITY_ARN, + AWS_STEPFUNCTIONS_STATEMACHINE_ARN, ) from amazon.opentelemetry.distro.patches._bedrock_patches import ( # noqa # pylint: disable=unused-import _BedrockAgentExtension, @@ -35,6 +37,20 @@ def _apply_botocore_instrumentation_patches() -> None: _apply_botocore_bedrock_patch() _apply_botocore_secretsmanager_patch() _apply_botocore_sns_patch() + _apply_botocore_stepfunctions_patch() + + +def _apply_botocore_stepfunctions_patch() -> None: + """Botocore instrumentation patch for StepFunctions + + This patch adds an extension to the upstream's list of known extensions for + StepFunctions. Extensions allow for custom logic for adding service-specific + information to spans, such as attributes. Specifically, we are adding logic + to add the `aws.stepfunctions.state_machine.arn` and `aws.stepfunctions.activity.arn` + attributes, to be used to generate RemoteTarget and achieve partity with the + Java instrumentation. + """ + _KNOWN_EXTENSIONS["stepfunctions"] = _lazy_load(".", "_StepFunctionsExtension") def _apply_botocore_sns_patch() -> None: @@ -149,6 +165,16 @@ def loader(): # END The OpenTelemetry Authors code +class _StepFunctionsExtension(_AwsSdkExtension): + def extract_attributes(self, attributes: _AttributeMapT): + state_machine_arn = self._call_context.params.get("stateMachineArn") + if state_machine_arn: + attributes[AWS_STEPFUNCTIONS_STATEMACHINE_ARN] = state_machine_arn + activity_arn = self._call_context.params.get("activityArn") + if activity_arn: + attributes[AWS_STEPFUNCTIONS_ACTIVITY_ARN] = activity_arn + + class _SecretsManagerExtension(_AwsSdkExtension): def extract_attributes(self, attributes: _AttributeMapT): """ 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 3dcec10db..92ebdb016 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 @@ -26,6 +26,8 @@ AWS_SPAN_KIND, AWS_SQS_QUEUE_NAME, AWS_SQS_QUEUE_URL, + AWS_STEPFUNCTIONS_ACTIVITY_ARN, + AWS_STEPFUNCTIONS_STATEMACHINE_ARN, ) from amazon.opentelemetry.distro._aws_metric_attribute_generator import _AwsMetricAttributeGenerator from amazon.opentelemetry.distro._aws_span_processing_util import GEN_AI_REQUEST_MODEL @@ -881,6 +883,7 @@ def test_normalize_remote_service_name_aws_sdk(self): self.validate_aws_sdk_service_normalization("Bedrock Runtime", "AWS::BedrockRuntime") self.validate_aws_sdk_service_normalization("Secrets Manager", "AWS::SecretsManager") self.validate_aws_sdk_service_normalization("SNS", "AWS::SNS") + self.validate_aws_sdk_service_normalization("SFN", "AWS::StepFunctions") def validate_aws_sdk_service_normalization(self, service_name: str, expected_remote_service: str): self._mock_attribute([SpanAttributes.RPC_SYSTEM, SpanAttributes.RPC_SERVICE], ["aws-api", service_name]) @@ -1114,6 +1117,31 @@ def test_sdk_client_span_with_remote_resource_attributes(self): self._validate_remote_resource_attributes("AWS::SNS::Topic", "arn:aws:sns:us-west-2:012345678901:test_topic") self._mock_attribute([AWS_SNS_TOPIC_ARN], [None]) + # Validate behaviour of AWS_STEPFUNCTIONS_STATEMACHINE_ARN attribute, then remove it. + self._mock_attribute( + [AWS_STEPFUNCTIONS_STATEMACHINE_ARN], + ["arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine"], + keys, + values + ) + self._validate_remote_resource_attributes( + "AWS::StepFunctions::StateMachine", + "arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine" + ) + self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None]) + + # Validate behaviour of AWS_STEPFUNCTIONS_ACTIVITY_ARN attribute, then remove it. + self._mock_attribute( + [AWS_STEPFUNCTIONS_ACTIVITY_ARN], + ["arn:aws:states:us-east-1:007003123456789012:activity:testActivity"], + keys, + values, + ) + self._validate_remote_resource_attributes( + "AWS::StepFunctions::Activity", "arn:aws:states:us-east-1:007003123456789012:activity:testActivity" + ) + self._mock_attribute([AWS_STEPFUNCTIONS_ACTIVITY_ARN], [None]) + self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None]) def test_client_db_span_with_remote_resource_attributes(self): 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 25b191670..048d02f17 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 @@ -28,6 +28,8 @@ _GEN_AI_REQUEST_MODEL: str = "genAiReuqestModelId" _SECRET_ARN: str = "arn:aws:secretsmanager:us-west-2:000000000000:secret:testSecret-ABCDEF" _TOPIC_ARN: str = "topicArn" +_STATE_MACHINE_ARN: str = "arn:aws:states:us-west-2:000000000000:stateMachine:testStateMachine" +_ACTIVITY_ARN: str = "arn:aws:states:us-east-1:007003123456789012:activity:testActivity" # Patch names GET_DISTRIBUTION_PATCH: str = ( @@ -149,6 +151,9 @@ def _test_unpatched_botocore_instrumentation(self): # SNS self.assertTrue("sns" in _KNOWN_EXTENSIONS, "Upstream has removed the SNS extension") + # StepFunctions + self.assertFalse("stepfunctions" in _KNOWN_EXTENSIONS, "Upstream has added a StepFunctions extension") + def _test_unpatched_gevent_instrumentation(self): self.assertFalse(gevent.monkey.is_module_patched("os"), "gevent os module has been patched") self.assertFalse(gevent.monkey.is_module_patched("thread"), "gevent thread module has been patched") @@ -223,6 +228,14 @@ def _test_patched_botocore_instrumentation(self): self.assertTrue("aws.sns.topic.arn" in sns_attributes) self.assertEqual(sns_attributes["aws.sns.topic.arn"], _TOPIC_ARN) + # StepFunctions + self.assertTrue("stepfunctions" in _KNOWN_EXTENSIONS) + stepfunctions_attributes: Dict[str, str] = _do_extract_stepfunctions_attributes() + self.assertTrue("aws.stepfunctions.state_machine.arn" in stepfunctions_attributes) + self.assertEqual(stepfunctions_attributes["aws.stepfunctions.state_machine.arn"], _STATE_MACHINE_ARN) + self.assertTrue("aws.stepfunctions.activity.arn" in stepfunctions_attributes) + self.assertEqual(stepfunctions_attributes["aws.stepfunctions.activity.arn"], _ACTIVITY_ARN) + def _test_patched_gevent_os_ssl_instrumentation(self): # Only ssl and os module should have been patched since the environment variable was set to 'os, ssl' self.assertTrue(gevent.monkey.is_module_patched("ssl"), "gevent ssl module has not been patched") @@ -399,6 +412,12 @@ def _do_extract_sns_attributes() -> Dict[str, str]: return _do_extract_attributes(service_name, params) +def _do_extract_stepfunctions_attributes() -> Dict[str, str]: + service_name: str = "stepfunctions" + params: Dict[str, str] = {"stateMachineArn": _STATE_MACHINE_ARN, "activityArn": _ACTIVITY_ARN} + return _do_extract_attributes(service_name, params) + + def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]: mock_call_context: MagicMock = MagicMock() mock_call_context.params = params From 525f0ef6997016dfa4eb521a02fe530f3783f9f0 Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Wed, 25 Sep 2024 14:23:23 -0700 Subject: [PATCH 04/11] feat: Add auto-instrumentation support for Kinesis Stream Consumer --- .../distro/_aws_attribute_keys.py | 1 + .../distro/_aws_metric_attribute_generator.py | 4 +++ .../distro/patches/_botocore_patches.py | 4 +++ .../test_aws_metric_attribute_generator.py | 25 +++++++++++++++++++ .../distro/test_instrumentation_patch.py | 5 +++- 5 files changed, 38 insertions(+), 1 deletion(-) 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 2d865ea1e..91d143b37 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 @@ -19,6 +19,7 @@ AWS_SQS_QUEUE_URL: str = "aws.sqs.queue.url" AWS_SQS_QUEUE_NAME: str = "aws.sqs.queue.name" AWS_KINESIS_STREAM_NAME: str = "aws.kinesis.stream.name" +AWS_KINESIS_STREAM_CONSUMERNAME: str = "aws.kinesis.stream.consumer_name" AWS_BEDROCK_DATA_SOURCE_ID: str = "aws.bedrock.data_source.id" AWS_BEDROCK_KNOWLEDGE_BASE_ID: str = "aws.bedrock.knowledge_base.id" AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id" 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 28a6cc116..297cd07df 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 @@ -11,6 +11,7 @@ AWS_BEDROCK_GUARDRAIL_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID, AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER, + AWS_KINESIS_STREAM_CONSUMERNAME, AWS_KINESIS_STREAM_NAME, AWS_LOCAL_OPERATION, AWS_LOCAL_SERVICE, @@ -393,6 +394,9 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri 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_CONSUMERNAME): + remote_resource_type = _NORMALIZED_KINESIS_SERVICE_NAME + "::StreamConsumer" + remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_KINESIS_STREAM_CONSUMERNAME)) 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)) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index 5a62d6002..3071cdb0f 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -4,6 +4,7 @@ import importlib from amazon.opentelemetry.distro._aws_attribute_keys import ( + AWS_KINESIS_STREAM_CONSUMERNAME, AWS_KINESIS_STREAM_NAME, AWS_SECRETSMANAGER_SECRET_ARN, AWS_SNS_TOPIC_ARN, @@ -204,3 +205,6 @@ def extract_attributes(self, attributes: _AttributeMapT): stream_name = self._call_context.params.get("StreamName") if stream_name: attributes[AWS_KINESIS_STREAM_NAME] = stream_name + consumer_name = self._call_context.params.get("ConsumerName") + if consumer_name: + attributes[AWS_KINESIS_STREAM_CONSUMERNAME] = consumer_name 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 92ebdb016..0a20de03e 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 @@ -13,6 +13,7 @@ AWS_BEDROCK_GUARDRAIL_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID, AWS_CONSUMER_PARENT_SPAN_KIND, + AWS_KINESIS_STREAM_CONSUMERNAME, AWS_KINESIS_STREAM_NAME, AWS_LOCAL_OPERATION, AWS_LOCAL_SERVICE, @@ -1142,6 +1143,30 @@ def test_sdk_client_span_with_remote_resource_attributes(self): ) self._mock_attribute([AWS_STEPFUNCTIONS_ACTIVITY_ARN], [None]) + # Validate behaviour of AWS_KINESIS_STREAM_CONSUMERNAME present, then remove it. + self._mock_attribute( + [AWS_KINESIS_STREAM_CONSUMERNAME], + ["aws_stream_consumername"], + keys, + values, + ) + self._validate_remote_resource_attributes( + "AWS::Kinesis::StreamConsumer", + "aws_stream_consumername" + ) + self._mock_attribute([AWS_KINESIS_STREAM_CONSUMERNAME], [None]) + + # Validate behaviour with both AWS_KINESIS_STREAM_NAME and AWS_KINESIS_STREAM_CONSUMERNAME + # present, then remove it. + self._mock_attribute( + [AWS_KINESIS_STREAM_NAME, AWS_KINESIS_STREAM_CONSUMERNAME], + ["aws_stream_name", "aws_stream_consumername"], + keys, + values, + ) + self._validate_remote_resource_attributes("AWS::Kinesis::Stream", "aws_stream_name") + self._mock_attribute([AWS_KINESIS_STREAM_NAME, AWS_KINESIS_STREAM_CONSUMERNAME], [None, None]) + self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None]) def test_client_db_span_with_remote_resource_attributes(self): 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 048d02f17..7f80f18e5 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 @@ -17,6 +17,7 @@ from opentelemetry.trace.span import Span _STREAM_NAME: str = "streamName" +_CONSUMER_NAME: str = "consumerName" _BUCKET_NAME: str = "bucketName" _QUEUE_NAME: str = "queueName" _QUEUE_URL: str = "queueUrl" @@ -175,6 +176,8 @@ def _test_patched_botocore_instrumentation(self): kinesis_attributes: Dict[str, str] = _do_extract_kinesis_attributes() self.assertTrue("aws.kinesis.stream.name" in kinesis_attributes) self.assertEqual(kinesis_attributes["aws.kinesis.stream.name"], _STREAM_NAME) + self.assertTrue("aws.kinesis.stream.consumer_name" in kinesis_attributes) + self.assertEqual(kinesis_attributes["aws.kinesis.stream.consumer_name"], _CONSUMER_NAME) # S3 self.assertTrue("s3" in _KNOWN_EXTENSIONS) @@ -356,7 +359,7 @@ def _reset_mocks(self): def _do_extract_kinesis_attributes() -> Dict[str, str]: service_name: str = "kinesis" - params: Dict[str, str] = {"StreamName": _STREAM_NAME} + params: Dict[str, str] = {"StreamName": _STREAM_NAME, "ConsumerName": _CONSUMER_NAME} return _do_extract_attributes(service_name, params) From 6f0b98e8846e9e18f96d8033647f6b81b5acb948 Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Wed, 25 Sep 2024 16:04:01 -0700 Subject: [PATCH 05/11] feat: Add auto-instrumentation support for Lambda --- .../distro/_aws_attribute_keys.py | 3 ++ .../distro/_aws_metric_attribute_generator.py | 11 +++++ .../distro/patches/_botocore_patches.py | 42 +++++++++++++++++++ .../test_aws_metric_attribute_generator.py | 23 ++++++++++ .../distro/test_instrumentation_patch.py | 19 +++++++++ 5 files changed, 98 insertions(+) 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 91d143b37..8361bec6d 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 @@ -28,3 +28,6 @@ AWS_SNS_TOPIC_ARN: str = "aws.sns.topic.arn" AWS_STEPFUNCTIONS_STATEMACHINE_ARN: str = "aws.stepfunctions.state_machine.arn" AWS_STEPFUNCTIONS_ACTIVITY_ARN: str = "aws.stepfunctions.activity.arn" +AWS_LAMBDA_FUNCTION_NAME: str = "aws.lambda.function.name" +AWS_LAMBDA_RESOURCEMAPPING_ID: str = "aws.lambda.resource_mapping.id" +AWS_LAMBDA_FUNCTION_ARN: str = "aws.lambda.function.arn" 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 297cd07df..55c578905 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 @@ -13,6 +13,9 @@ AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER, AWS_KINESIS_STREAM_CONSUMERNAME, AWS_KINESIS_STREAM_NAME, + AWS_LAMBDA_FUNCTION_ARN, + AWS_LAMBDA_FUNCTION_NAME, + AWS_LAMBDA_RESOURCEMAPPING_ID, AWS_LOCAL_OPERATION, AWS_LOCAL_SERVICE, AWS_REMOTE_DB_USER, @@ -96,6 +99,7 @@ _NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager" _NORMALIZED_SNS_SERVICE_NAME: str = "AWS::SNS" _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME: str = "AWS::StepFunctions" +_NORMALIZED_LAMBDA_SERVICE_NAME: str = "AWS::Lambda" _DB_CONNECTION_STRING_TYPE: str = "DB::Connection" # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present. @@ -442,6 +446,13 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri 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)) + elif is_key_present(span, AWS_LAMBDA_RESOURCEMAPPING_ID): + remote_resource_type = _NORMALIZED_LAMBDA_SERVICE_NAME + "::EventSourceMapping" + remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_RESOURCEMAPPING_ID)) + elif is_key_present(span, AWS_LAMBDA_FUNCTION_NAME): + remote_resource_type = _NORMALIZED_LAMBDA_SERVICE_NAME + "::Function" + remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_FUNCTION_NAME)) + cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_FUNCTION_ARN)) elif is_db_span(span): remote_resource_type = _DB_CONNECTION_STRING_TYPE remote_resource_identifier = _get_db_connection(span) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index 3071cdb0f..52178646a 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -6,6 +6,9 @@ from amazon.opentelemetry.distro._aws_attribute_keys import ( AWS_KINESIS_STREAM_CONSUMERNAME, AWS_KINESIS_STREAM_NAME, + AWS_LAMBDA_FUNCTION_ARN, + AWS_LAMBDA_FUNCTION_NAME, + AWS_LAMBDA_RESOURCEMAPPING_ID, AWS_SECRETSMANAGER_SECRET_ARN, AWS_SNS_TOPIC_ARN, AWS_SQS_QUEUE_NAME, @@ -20,6 +23,7 @@ _BedrockRuntimeExtension, ) from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS +from opentelemetry.instrumentation.botocore.extensions.lmbd import _LambdaExtension from opentelemetry.instrumentation.botocore.extensions.sns import _SnsExtension from opentelemetry.instrumentation.botocore.extensions.sqs import _SqsExtension from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension, _BotoResultT @@ -39,6 +43,44 @@ def _apply_botocore_instrumentation_patches() -> None: _apply_botocore_secretsmanager_patch() _apply_botocore_sns_patch() _apply_botocore_stepfunctions_patch() + _apply_botocore_lambda_patch() + + +def _apply_botocore_lambda_patch() -> None: + """Botocore instrumentation patch for Lambda + + This patch adds an extension to the upstream's list of known extensions for Lambda. + Extensions allow for custom logic for adding service-specific information to spans, + such as attributes. Specifically, we are adding logic to add the + `aws.lambda.function.name` and `aws.lambda.resource_mapping.id` attributes + + Sidenote: There exists SpanAttributes.FAAS_INVOKED_NAME for invoke operations + in upstream. However, we want to cover more operations to extract 'FunctionName', + so we define `aws.lambda.function.name` separately. Additionally, this helps + us maintain naming consistency with the other AWS resources. + """ + old_extract_attributes = _LambdaExtension.extract_attributes + + def patch_extract_attributes(self, attributes: _AttributeMapT): + old_extract_attributes(self, attributes) + function_name = self._call_context.params.get("FunctionName") + if function_name: + attributes[AWS_LAMBDA_FUNCTION_NAME] = function_name + resource_mapping_id = self._call_context.params.get("UUID") + if resource_mapping_id: + attributes[AWS_LAMBDA_RESOURCEMAPPING_ID] = resource_mapping_id + + old_on_success = _LambdaExtension.on_success + + def patch_on_success(self, span: Span, result: _BotoResultT): + old_on_success(self, span, result) + lambda_configuration = result.get("Configuration", {}) + function_arn = lambda_configuration.get("FunctionArn") + if function_arn: + span.set_attribute(AWS_LAMBDA_FUNCTION_ARN, function_arn) + + _LambdaExtension.extract_attributes = patch_extract_attributes + _LambdaExtension.on_success = patch_on_success def _apply_botocore_stepfunctions_patch() -> None: 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 0a20de03e..73d201c11 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 @@ -15,6 +15,8 @@ AWS_CONSUMER_PARENT_SPAN_KIND, AWS_KINESIS_STREAM_CONSUMERNAME, AWS_KINESIS_STREAM_NAME, + AWS_LAMBDA_FUNCTION_NAME, + AWS_LAMBDA_RESOURCEMAPPING_ID, AWS_LOCAL_OPERATION, AWS_LOCAL_SERVICE, AWS_REMOTE_DB_USER, @@ -1167,6 +1169,27 @@ def test_sdk_client_span_with_remote_resource_attributes(self): self._validate_remote_resource_attributes("AWS::Kinesis::Stream", "aws_stream_name") self._mock_attribute([AWS_KINESIS_STREAM_NAME, AWS_KINESIS_STREAM_CONSUMERNAME], [None, None]) + # Validate behaviour of AWS_LAMBDA_FUNCTION_NAME attribute, then remove it. + self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME], ["aws_lambda_function_name"], keys, values) + self._validate_remote_resource_attributes("AWS::Lambda::Function", "aws_lambda_function_name") + self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME], [None]) + + # Validate behaviour of AWS_LAMBDA_RESOURCEMAPPING_ID attribute, then remove it. + self._mock_attribute([AWS_LAMBDA_RESOURCEMAPPING_ID], ["aws_event_source_mapping_id"], keys, values) + self._validate_remote_resource_attributes("AWS::Lambda::EventSourceMapping", "aws_event_source_mapping_id") + self._mock_attribute([AWS_LAMBDA_RESOURCEMAPPING_ID], [None]) + + # Validate behaviour of both AWS_LAMBDA_FUNCTION_NAME and AWS_LAMBDA_RESOURCE_MAPPING_ID, + # then remove it. + self._mock_attribute( + [AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_RESOURCEMAPPING_ID], + ["aws_lambda_function_name", "aws_event_source_mapping_id"], + keys, + values + ) + self._validate_remote_resource_attributes("AWS::Lambda::EventSourceMapping", "aws_event_source_mapping_id") + self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_RESOURCEMAPPING_ID], [None, None]) + self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None]) def test_client_db_span_with_remote_resource_attributes(self): 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 7f80f18e5..5e55e2f55 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 @@ -31,6 +31,8 @@ _TOPIC_ARN: str = "topicArn" _STATE_MACHINE_ARN: str = "arn:aws:states:us-west-2:000000000000:stateMachine:testStateMachine" _ACTIVITY_ARN: str = "arn:aws:states:us-east-1:007003123456789012:activity:testActivity" +_LAMBDA_FUNCTION_NAME: str = "lambdaFunctionName" +_LAMBDA_SOURCE_MAPPING_ID: str = "lambdaEventSourceMappingID" # Patch names GET_DISTRIBUTION_PATCH: str = ( @@ -155,6 +157,9 @@ def _test_unpatched_botocore_instrumentation(self): # StepFunctions self.assertFalse("stepfunctions" in _KNOWN_EXTENSIONS, "Upstream has added a StepFunctions extension") + # Lambda + self.assertTrue("lambda" in _KNOWN_EXTENSIONS, "Upstream has removed the Lambda extension") + def _test_unpatched_gevent_instrumentation(self): self.assertFalse(gevent.monkey.is_module_patched("os"), "gevent os module has been patched") self.assertFalse(gevent.monkey.is_module_patched("thread"), "gevent thread module has been patched") @@ -239,6 +244,14 @@ def _test_patched_botocore_instrumentation(self): self.assertTrue("aws.stepfunctions.activity.arn" in stepfunctions_attributes) self.assertEqual(stepfunctions_attributes["aws.stepfunctions.activity.arn"], _ACTIVITY_ARN) + # Lambda + self.assertTrue("lambda" in _KNOWN_EXTENSIONS) + lambda_attributes: Dict[str, str] = _do_extract_lambda_attributes() + self.assertTrue("aws.lambda.function.name" in lambda_attributes) + self.assertEqual(lambda_attributes["aws.lambda.function.name"], _LAMBDA_FUNCTION_NAME) + self.assertTrue("aws.lambda.resource_mapping.id" in lambda_attributes) + self.assertEqual(lambda_attributes["aws.lambda.resource_mapping.id"], _LAMBDA_SOURCE_MAPPING_ID) + def _test_patched_gevent_os_ssl_instrumentation(self): # Only ssl and os module should have been patched since the environment variable was set to 'os, ssl' self.assertTrue(gevent.monkey.is_module_patched("ssl"), "gevent ssl module has not been patched") @@ -421,6 +434,12 @@ def _do_extract_stepfunctions_attributes() -> Dict[str, str]: return _do_extract_attributes(service_name, params) +def _do_extract_lambda_attributes() -> Dict[str, str]: + service_name: str = "lambda" + params: Dict[str, str] = {"FunctionName": _LAMBDA_FUNCTION_NAME, "UUID": _LAMBDA_SOURCE_MAPPING_ID} + return _do_extract_attributes(service_name, params) + + def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]: mock_call_context: MagicMock = MagicMock() mock_call_context.params = params From 4df8bbd871d809de52086db3f699aa44ad7aeb1b Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Wed, 25 Sep 2024 16:06:48 -0700 Subject: [PATCH 06/11] chore: Apply code style changes --- .../distro/patches/_botocore_patches.py | 11 +++++----- .../test_aws_metric_attribute_generator.py | 20 ++++++++----------- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index 52178646a..dc3c3dd4f 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -51,7 +51,7 @@ def _apply_botocore_lambda_patch() -> None: This patch adds an extension to the upstream's list of known extensions for Lambda. Extensions allow for custom logic for adding service-specific information to spans, - such as attributes. Specifically, we are adding logic to add the + such as attributes. Specifically, we are adding logic to add the `aws.lambda.function.name` and `aws.lambda.resource_mapping.id` attributes Sidenote: There exists SpanAttributes.FAAS_INVOKED_NAME for invoke operations @@ -70,7 +70,7 @@ def patch_extract_attributes(self, attributes: _AttributeMapT): if resource_mapping_id: attributes[AWS_LAMBDA_RESOURCEMAPPING_ID] = resource_mapping_id - old_on_success = _LambdaExtension.on_success + old_on_success = _LambdaExtension.on_success def patch_on_success(self, span: Span, result: _BotoResultT): old_on_success(self, span, result) @@ -86,7 +86,7 @@ def patch_on_success(self, span: Span, result: _BotoResultT): def _apply_botocore_stepfunctions_patch() -> None: """Botocore instrumentation patch for StepFunctions - This patch adds an extension to the upstream's list of known extensions for + This patch adds an extension to the upstream's list of known extensions for StepFunctions. Extensions allow for custom logic for adding service-specific information to spans, such as attributes. Specifically, we are adding logic to add the `aws.stepfunctions.state_machine.arn` and `aws.stepfunctions.activity.arn` @@ -101,7 +101,7 @@ def _apply_botocore_sns_patch() -> None: This patch adds an extension to the upstream's list of known extensions for SNS. Extensions allow for custom logic for adding service-specific information to - spans, such as attributes. Specifically, we are adding logic to add the + spans, such as attributes. Specifically, we are adding logic to add the `aws.sns.topic.arn` attribute, to be used to generate RemoteTarget and achieve parity with the Java instrumentation. @@ -119,12 +119,13 @@ def patch_extract_attributes(self, attributes: _AttributeMapT): _SnsExtension.extract_attributes = patch_extract_attributes + def _apply_botocore_secretsmanager_patch() -> None: """Botocore instrumentation patch for SecretsManager This patch adds an extension to the upstream's list of known extension for SecretsManager. Extensions allow for custom logic for adding service-specific information to spans, such as - attributes. Specifically, we are adding logic to add the `aws.secretsmanager.secret.arn` + attributes. Specifically, we are adding logic to add the `aws.secretsmanager.secret.arn` attribute, to be used to generate RemoteTarget and achieve parity with the Java instrumentation. """ 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 73d201c11..97c0caf50 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 @@ -1108,7 +1108,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self): [AWS_SECRETSMANAGER_SECRET_ARN], ["arn:aws:secretsmanager:us-east-1:123456789012:secret:secret_name-lERW9H"], keys, - values + values, ) self._validate_remote_resource_attributes( "AWS::SecretsManager::Secret", "arn:aws:secretsmanager:us-east-1:123456789012:secret:secret_name-lERW9H" @@ -1122,14 +1122,13 @@ def test_sdk_client_span_with_remote_resource_attributes(self): # Validate behaviour of AWS_STEPFUNCTIONS_STATEMACHINE_ARN attribute, then remove it. self._mock_attribute( - [AWS_STEPFUNCTIONS_STATEMACHINE_ARN], - ["arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine"], - keys, - values + [AWS_STEPFUNCTIONS_STATEMACHINE_ARN], + ["arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine"], + keys, + values, ) self._validate_remote_resource_attributes( - "AWS::StepFunctions::StateMachine", - "arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine" + "AWS::StepFunctions::StateMachine", "arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine" ) self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None]) @@ -1152,10 +1151,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self): keys, values, ) - self._validate_remote_resource_attributes( - "AWS::Kinesis::StreamConsumer", - "aws_stream_consumername" - ) + self._validate_remote_resource_attributes("AWS::Kinesis::StreamConsumer", "aws_stream_consumername") self._mock_attribute([AWS_KINESIS_STREAM_CONSUMERNAME], [None]) # Validate behaviour with both AWS_KINESIS_STREAM_NAME and AWS_KINESIS_STREAM_CONSUMERNAME @@ -1185,7 +1181,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self): [AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_RESOURCEMAPPING_ID], ["aws_lambda_function_name", "aws_event_source_mapping_id"], keys, - values + values, ) self._validate_remote_resource_attributes("AWS::Lambda::EventSourceMapping", "aws_event_source_mapping_id") self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_RESOURCEMAPPING_ID], [None, None]) From 6765867fecff330bc91aa6671c4e932404ddee96 Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Wed, 25 Sep 2024 16:44:58 -0700 Subject: [PATCH 07/11] chore: Manually fix linter errors --- .../opentelemetry/distro/_aws_metric_attribute_generator.py | 2 +- .../amazon/opentelemetry/distro/patches/_botocore_patches.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) 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 55c578905..9758a8612 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 @@ -374,7 +374,7 @@ def _generate_remote_operation(span: ReadableSpan) -> str: return remote_operation -# pylint: disable=too-many-branches +# pylint: disable=too-many-branches,too-many-statements def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttributes) -> None: """ Remote resource attributes {@link AwsAttributeKeys#AWS_REMOTE_RESOURCE_TYPE} and {@link diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index dc3c3dd4f..c33fea56c 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -230,6 +230,7 @@ def extract_attributes(self, attributes: _AttributeMapT): if secret_id and secret_id.startswith("arn:aws:secretsmanager:"): attributes[AWS_SECRETSMANAGER_SECRET_ARN] = secret_id + # pylint: disable=no-self-use def on_success(self, span: Span, result: _BotoResultT): secret_arn = result.get("ARN") if secret_arn: From f05d295378db9e40e24b49bbd08b15b21eaae05f Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Thu, 26 Sep 2024 11:23:45 -0700 Subject: [PATCH 08/11] fix: Correct remote_resource_identifier attribute values for SNS, SecretsManager, and SFN --- .../distro/_aws_metric_attribute_generator.py | 12 ++++++++---- .../distro/test_aws_metric_attribute_generator.py | 8 ++++---- 2 files changed, 12 insertions(+), 8 deletions(-) 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 9758a8612..dbc1af6a5 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 @@ -436,16 +436,20 @@ 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)) + remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_SECRETSMANAGER_SECRET_ARN)).split(":")[-1] + 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)) + remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_SNS_TOPIC_ARN)).split(":")[-1] + 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)) + remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_STATEMACHINE_ARN)).split(":")[-1] + 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)) + remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN)).split(":")[-1] + cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN)) elif is_key_present(span, AWS_LAMBDA_RESOURCEMAPPING_ID): remote_resource_type = _NORMALIZED_LAMBDA_SERVICE_NAME + "::EventSourceMapping" remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_RESOURCEMAPPING_ID)) 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 97c0caf50..dc82b3a90 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 @@ -1111,13 +1111,13 @@ def test_sdk_client_span_with_remote_resource_attributes(self): values, ) self._validate_remote_resource_attributes( - "AWS::SecretsManager::Secret", "arn:aws:secretsmanager:us-east-1:123456789012:secret:secret_name-lERW9H" + "AWS::SecretsManager::Secret", "secret_name-lERW9H" ) self._mock_attribute([AWS_SECRETSMANAGER_SECRET_ARN], [None]) # Validate behaviour of AWS_SNS_TOPIC_ARN attribute, then remove it. self._mock_attribute([AWS_SNS_TOPIC_ARN], ["arn:aws:sns:us-west-2:012345678901:test_topic"], keys, values) - self._validate_remote_resource_attributes("AWS::SNS::Topic", "arn:aws:sns:us-west-2:012345678901:test_topic") + self._validate_remote_resource_attributes("AWS::SNS::Topic", "test_topic") self._mock_attribute([AWS_SNS_TOPIC_ARN], [None]) # Validate behaviour of AWS_STEPFUNCTIONS_STATEMACHINE_ARN attribute, then remove it. @@ -1128,7 +1128,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self): values, ) self._validate_remote_resource_attributes( - "AWS::StepFunctions::StateMachine", "arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine" + "AWS::StepFunctions::StateMachine", "test_state_machine" ) self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None]) @@ -1140,7 +1140,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self): values, ) self._validate_remote_resource_attributes( - "AWS::StepFunctions::Activity", "arn:aws:states:us-east-1:007003123456789012:activity:testActivity" + "AWS::StepFunctions::Activity", "testActivity" ) self._mock_attribute([AWS_STEPFUNCTIONS_ACTIVITY_ARN], [None]) From b0c266c598eb5a08c16a7903ee21086120a61173 Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Thu, 26 Sep 2024 11:32:36 -0700 Subject: [PATCH 09/11] clean: Remove auto-instrumentation support for Kinesis Stream Consumer --- .../distro/_aws_attribute_keys.py | 1 - .../distro/_aws_metric_attribute_generator.py | 4 ---- .../distro/patches/_botocore_patches.py | 4 ---- .../test_aws_metric_attribute_generator.py | 22 ------------------- .../distro/test_instrumentation_patch.py | 5 +---- 5 files changed, 1 insertion(+), 35 deletions(-) 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 8361bec6d..855b9c42a 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 @@ -19,7 +19,6 @@ AWS_SQS_QUEUE_URL: str = "aws.sqs.queue.url" AWS_SQS_QUEUE_NAME: str = "aws.sqs.queue.name" AWS_KINESIS_STREAM_NAME: str = "aws.kinesis.stream.name" -AWS_KINESIS_STREAM_CONSUMERNAME: str = "aws.kinesis.stream.consumer_name" AWS_BEDROCK_DATA_SOURCE_ID: str = "aws.bedrock.data_source.id" AWS_BEDROCK_KNOWLEDGE_BASE_ID: str = "aws.bedrock.knowledge_base.id" AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id" 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 dbc1af6a5..7f43ca1f2 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 @@ -11,7 +11,6 @@ AWS_BEDROCK_GUARDRAIL_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID, AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER, - AWS_KINESIS_STREAM_CONSUMERNAME, AWS_KINESIS_STREAM_NAME, AWS_LAMBDA_FUNCTION_ARN, AWS_LAMBDA_FUNCTION_NAME, @@ -398,9 +397,6 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri 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_CONSUMERNAME): - remote_resource_type = _NORMALIZED_KINESIS_SERVICE_NAME + "::StreamConsumer" - remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_KINESIS_STREAM_CONSUMERNAME)) 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)) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index c33fea56c..07dc8a0d5 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -4,7 +4,6 @@ import importlib from amazon.opentelemetry.distro._aws_attribute_keys import ( - AWS_KINESIS_STREAM_CONSUMERNAME, AWS_KINESIS_STREAM_NAME, AWS_LAMBDA_FUNCTION_ARN, AWS_LAMBDA_FUNCTION_NAME, @@ -249,6 +248,3 @@ def extract_attributes(self, attributes: _AttributeMapT): stream_name = self._call_context.params.get("StreamName") if stream_name: attributes[AWS_KINESIS_STREAM_NAME] = stream_name - consumer_name = self._call_context.params.get("ConsumerName") - if consumer_name: - attributes[AWS_KINESIS_STREAM_CONSUMERNAME] = consumer_name 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 dc82b3a90..4dc235236 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 @@ -13,7 +13,6 @@ AWS_BEDROCK_GUARDRAIL_ID, AWS_BEDROCK_KNOWLEDGE_BASE_ID, AWS_CONSUMER_PARENT_SPAN_KIND, - AWS_KINESIS_STREAM_CONSUMERNAME, AWS_KINESIS_STREAM_NAME, AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_RESOURCEMAPPING_ID, @@ -1144,27 +1143,6 @@ def test_sdk_client_span_with_remote_resource_attributes(self): ) self._mock_attribute([AWS_STEPFUNCTIONS_ACTIVITY_ARN], [None]) - # Validate behaviour of AWS_KINESIS_STREAM_CONSUMERNAME present, then remove it. - self._mock_attribute( - [AWS_KINESIS_STREAM_CONSUMERNAME], - ["aws_stream_consumername"], - keys, - values, - ) - self._validate_remote_resource_attributes("AWS::Kinesis::StreamConsumer", "aws_stream_consumername") - self._mock_attribute([AWS_KINESIS_STREAM_CONSUMERNAME], [None]) - - # Validate behaviour with both AWS_KINESIS_STREAM_NAME and AWS_KINESIS_STREAM_CONSUMERNAME - # present, then remove it. - self._mock_attribute( - [AWS_KINESIS_STREAM_NAME, AWS_KINESIS_STREAM_CONSUMERNAME], - ["aws_stream_name", "aws_stream_consumername"], - keys, - values, - ) - self._validate_remote_resource_attributes("AWS::Kinesis::Stream", "aws_stream_name") - self._mock_attribute([AWS_KINESIS_STREAM_NAME, AWS_KINESIS_STREAM_CONSUMERNAME], [None, None]) - # Validate behaviour of AWS_LAMBDA_FUNCTION_NAME attribute, then remove it. self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME], ["aws_lambda_function_name"], keys, values) self._validate_remote_resource_attributes("AWS::Lambda::Function", "aws_lambda_function_name") 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 5e55e2f55..b27d5e799 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 @@ -17,7 +17,6 @@ from opentelemetry.trace.span import Span _STREAM_NAME: str = "streamName" -_CONSUMER_NAME: str = "consumerName" _BUCKET_NAME: str = "bucketName" _QUEUE_NAME: str = "queueName" _QUEUE_URL: str = "queueUrl" @@ -181,8 +180,6 @@ def _test_patched_botocore_instrumentation(self): kinesis_attributes: Dict[str, str] = _do_extract_kinesis_attributes() self.assertTrue("aws.kinesis.stream.name" in kinesis_attributes) self.assertEqual(kinesis_attributes["aws.kinesis.stream.name"], _STREAM_NAME) - self.assertTrue("aws.kinesis.stream.consumer_name" in kinesis_attributes) - self.assertEqual(kinesis_attributes["aws.kinesis.stream.consumer_name"], _CONSUMER_NAME) # S3 self.assertTrue("s3" in _KNOWN_EXTENSIONS) @@ -372,7 +369,7 @@ def _reset_mocks(self): def _do_extract_kinesis_attributes() -> Dict[str, str]: service_name: str = "kinesis" - params: Dict[str, str] = {"StreamName": _STREAM_NAME, "ConsumerName": _CONSUMER_NAME} + params: Dict[str, str] = {"StreamName": _STREAM_NAME} return _do_extract_attributes(service_name, params) From d5e06667a56f043db66ba5b2be6690a4d06c470b Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Thu, 26 Sep 2024 11:36:19 -0700 Subject: [PATCH 10/11] chore: Apply code style changes --- .../distro/_aws_metric_attribute_generator.py | 16 ++++++++++++---- .../test_aws_metric_attribute_generator.py | 12 +++--------- 2 files changed, 15 insertions(+), 13 deletions(-) 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 7f43ca1f2..40a4fd46b 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 @@ -432,7 +432,9 @@ 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(span.attributes.get(AWS_SECRETSMANAGER_SECRET_ARN)).split( + ":" + )[-1] 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" @@ -440,11 +442,17 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri 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] - cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_STATEMACHINE_ARN)) + remote_resource_identifier = _escape_delimiters( + span.attributes.get(AWS_STEPFUNCTIONS_STATEMACHINE_ARN) + ).split(":")[-1] + 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(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN)).split( + ":" + )[-1] cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN)) elif is_key_present(span, AWS_LAMBDA_RESOURCEMAPPING_ID): remote_resource_type = _NORMALIZED_LAMBDA_SERVICE_NAME + "::EventSourceMapping" 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 4dc235236..f9da8bb9a 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 @@ -1109,9 +1109,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self): keys, values, ) - self._validate_remote_resource_attributes( - "AWS::SecretsManager::Secret", "secret_name-lERW9H" - ) + self._validate_remote_resource_attributes("AWS::SecretsManager::Secret", "secret_name-lERW9H") self._mock_attribute([AWS_SECRETSMANAGER_SECRET_ARN], [None]) # Validate behaviour of AWS_SNS_TOPIC_ARN attribute, then remove it. @@ -1126,9 +1124,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self): keys, values, ) - self._validate_remote_resource_attributes( - "AWS::StepFunctions::StateMachine", "test_state_machine" - ) + self._validate_remote_resource_attributes("AWS::StepFunctions::StateMachine", "test_state_machine") self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None]) # Validate behaviour of AWS_STEPFUNCTIONS_ACTIVITY_ARN attribute, then remove it. @@ -1138,9 +1134,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self): keys, values, ) - self._validate_remote_resource_attributes( - "AWS::StepFunctions::Activity", "testActivity" - ) + self._validate_remote_resource_attributes("AWS::StepFunctions::Activity", "testActivity") self._mock_attribute([AWS_STEPFUNCTIONS_ACTIVITY_ARN], [None]) # Validate behaviour of AWS_LAMBDA_FUNCTION_NAME attribute, then remove it. From d5dece2398778cd9bad7b605fa75d39dee86f30e Mon Sep 17 00:00:00 2001 From: yiyuanh Date: Thu, 26 Sep 2024 11:55:54 -0700 Subject: [PATCH 11/11] fix: Standardize Lambda Function Name Parameter --- .../opentelemetry/distro/patches/_botocore_patches.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index 07dc8a0d5..069e39c87 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -62,8 +62,12 @@ def _apply_botocore_lambda_patch() -> None: def patch_extract_attributes(self, attributes: _AttributeMapT): old_extract_attributes(self, attributes) - function_name = self._call_context.params.get("FunctionName") - if function_name: + # This param can be passed as an arn or a name. We standardize it to be the name. + function_name_param = self._call_context.params.get("FunctionName") + if function_name_param: + function_name = function_name_param + if function_name_param.startswith("arn:aws:lambda:"): + function_name = function_name_param.split(":")[-1] attributes[AWS_LAMBDA_FUNCTION_NAME] = function_name resource_mapping_id = self._call_context.params.get("UUID") if resource_mapping_id: