55
66from amazon .opentelemetry .distro ._aws_attribute_keys import (
77 AWS_KINESIS_STREAM_NAME ,
8+ AWS_LAMBDA_FUNCTION_ARN ,
9+ AWS_LAMBDA_FUNCTION_NAME ,
10+ AWS_LAMBDA_RESOURCEMAPPING_ID ,
11+ AWS_SECRETSMANAGER_SECRET_ARN ,
12+ AWS_SNS_TOPIC_ARN ,
813 AWS_SQS_QUEUE_NAME ,
914 AWS_SQS_QUEUE_URL ,
15+ AWS_STEPFUNCTIONS_ACTIVITY_ARN ,
16+ AWS_STEPFUNCTIONS_STATEMACHINE_ARN ,
1017)
1118from amazon .opentelemetry .distro .patches ._bedrock_patches import ( # noqa # pylint: disable=unused-import
1219 _BedrockAgentExtension ,
1522 _BedrockRuntimeExtension ,
1623)
1724from opentelemetry .instrumentation .botocore .extensions import _KNOWN_EXTENSIONS
25+ from opentelemetry .instrumentation .botocore .extensions .lmbd import _LambdaExtension
26+ from opentelemetry .instrumentation .botocore .extensions .sns import _SnsExtension
1827from opentelemetry .instrumentation .botocore .extensions .sqs import _SqsExtension
19- from opentelemetry .instrumentation .botocore .extensions .types import _AttributeMapT , _AwsSdkExtension
28+ from opentelemetry .instrumentation .botocore .extensions .types import _AttributeMapT , _AwsSdkExtension , _BotoResultT
2029from opentelemetry .semconv .trace import SpanAttributes
30+ from opentelemetry .trace .span import Span
2131
2232
2333def _apply_botocore_instrumentation_patches () -> None :
@@ -29,6 +39,100 @@ def _apply_botocore_instrumentation_patches() -> None:
2939 _apply_botocore_s3_patch ()
3040 _apply_botocore_sqs_patch ()
3141 _apply_botocore_bedrock_patch ()
42+ _apply_botocore_secretsmanager_patch ()
43+ _apply_botocore_sns_patch ()
44+ _apply_botocore_stepfunctions_patch ()
45+ _apply_botocore_lambda_patch ()
46+
47+
48+ def _apply_botocore_lambda_patch () -> None :
49+ """Botocore instrumentation patch for Lambda
50+
51+ This patch adds an extension to the upstream's list of known extensions for Lambda.
52+ Extensions allow for custom logic for adding service-specific information to spans,
53+ such as attributes. Specifically, we are adding logic to add the
54+ `aws.lambda.function.name` and `aws.lambda.resource_mapping.id` attributes
55+
56+ Sidenote: There exists SpanAttributes.FAAS_INVOKED_NAME for invoke operations
57+ in upstream. However, we want to cover more operations to extract 'FunctionName',
58+ so we define `aws.lambda.function.name` separately. Additionally, this helps
59+ us maintain naming consistency with the other AWS resources.
60+ """
61+ old_extract_attributes = _LambdaExtension .extract_attributes
62+
63+ def patch_extract_attributes (self , attributes : _AttributeMapT ):
64+ old_extract_attributes (self , attributes )
65+ # This param can be passed as an arn or a name. We standardize it to be the name.
66+ function_name_param = self ._call_context .params .get ("FunctionName" )
67+ if function_name_param :
68+ function_name = function_name_param
69+ if function_name_param .startswith ("arn:aws:lambda:" ):
70+ function_name = function_name_param .split (":" )[- 1 ]
71+ attributes [AWS_LAMBDA_FUNCTION_NAME ] = function_name
72+ resource_mapping_id = self ._call_context .params .get ("UUID" )
73+ if resource_mapping_id :
74+ attributes [AWS_LAMBDA_RESOURCEMAPPING_ID ] = resource_mapping_id
75+
76+ old_on_success = _LambdaExtension .on_success
77+
78+ def patch_on_success (self , span : Span , result : _BotoResultT ):
79+ old_on_success (self , span , result )
80+ lambda_configuration = result .get ("Configuration" , {})
81+ function_arn = lambda_configuration .get ("FunctionArn" )
82+ if function_arn :
83+ span .set_attribute (AWS_LAMBDA_FUNCTION_ARN , function_arn )
84+
85+ _LambdaExtension .extract_attributes = patch_extract_attributes
86+ _LambdaExtension .on_success = patch_on_success
87+
88+
89+ def _apply_botocore_stepfunctions_patch () -> None :
90+ """Botocore instrumentation patch for StepFunctions
91+
92+ This patch adds an extension to the upstream's list of known extensions for
93+ StepFunctions. Extensions allow for custom logic for adding service-specific
94+ information to spans, such as attributes. Specifically, we are adding logic
95+ to add the `aws.stepfunctions.state_machine.arn` and `aws.stepfunctions.activity.arn`
96+ attributes, to be used to generate RemoteTarget and achieve partity with the
97+ Java instrumentation.
98+ """
99+ _KNOWN_EXTENSIONS ["stepfunctions" ] = _lazy_load ("." , "_StepFunctionsExtension" )
100+
101+
102+ def _apply_botocore_sns_patch () -> None :
103+ """Botocore instrumentation patch for SNS
104+
105+ This patch adds an extension to the upstream's list of known extensions for SNS.
106+ Extensions allow for custom logic for adding service-specific information to
107+ spans, such as attributes. Specifically, we are adding logic to add the
108+ `aws.sns.topic.arn` attribute, to be used to generate RemoteTarget and achieve
109+ parity with the Java instrumentation.
110+
111+ Sidenote: There exists SpanAttributes.MESSAGING_DESTINATION_NAME in the upstream
112+ logic that we could re-purpose here. We do not use it here to maintain consistent
113+ naming patterns with other AWS resources.
114+ """
115+ old_extract_attributes = _SnsExtension .extract_attributes
116+
117+ def patch_extract_attributes (self , attributes : _AttributeMapT ):
118+ old_extract_attributes (self , attributes )
119+ topic_arn = self ._call_context .params .get ("TopicArn" )
120+ if topic_arn :
121+ attributes [AWS_SNS_TOPIC_ARN ] = topic_arn
122+
123+ _SnsExtension .extract_attributes = patch_extract_attributes
124+
125+
126+ def _apply_botocore_secretsmanager_patch () -> None :
127+ """Botocore instrumentation patch for SecretsManager
128+
129+ This patch adds an extension to the upstream's list of known extension for SecretsManager.
130+ Extensions allow for custom logic for adding service-specific information to spans, such as
131+ attributes. Specifically, we are adding logic to add the `aws.secretsmanager.secret.arn`
132+ attribute, to be used to generate RemoteTarget and achieve parity with the Java
133+ instrumentation.
134+ """
135+ _KNOWN_EXTENSIONS ["secretsmanager" ] = _lazy_load ("." , "_SecretsManagerExtension" )
32136
33137
34138def _apply_botocore_kinesis_patch () -> None :
@@ -108,6 +212,34 @@ def loader():
108212# END The OpenTelemetry Authors code
109213
110214
215+ class _StepFunctionsExtension (_AwsSdkExtension ):
216+ def extract_attributes (self , attributes : _AttributeMapT ):
217+ state_machine_arn = self ._call_context .params .get ("stateMachineArn" )
218+ if state_machine_arn :
219+ attributes [AWS_STEPFUNCTIONS_STATEMACHINE_ARN ] = state_machine_arn
220+ activity_arn = self ._call_context .params .get ("activityArn" )
221+ if activity_arn :
222+ attributes [AWS_STEPFUNCTIONS_ACTIVITY_ARN ] = activity_arn
223+
224+
225+ class _SecretsManagerExtension (_AwsSdkExtension ):
226+ def extract_attributes (self , attributes : _AttributeMapT ):
227+ """
228+ SecretId can be secret name or secret arn, the function extracts attributes
229+ only if the SecretId parameter is provided as an arn which starts with
230+ `arn:aws:secretsmanager:`
231+ """
232+ secret_id = self ._call_context .params .get ("SecretId" )
233+ if secret_id and secret_id .startswith ("arn:aws:secretsmanager:" ):
234+ attributes [AWS_SECRETSMANAGER_SECRET_ARN ] = secret_id
235+
236+ # pylint: disable=no-self-use
237+ def on_success (self , span : Span , result : _BotoResultT ):
238+ secret_arn = result .get ("ARN" )
239+ if secret_arn :
240+ span .set_attribute (AWS_SECRETSMANAGER_SECRET_ARN , secret_arn )
241+
242+
111243class _S3Extension (_AwsSdkExtension ):
112244 def extract_attributes (self , attributes : _AttributeMapT ):
113245 bucket_name = self ._call_context .params .get ("Bucket" )
0 commit comments