Skip to content

Commit a23151a

Browse files
blairhyy-amazonblairhyy-amazon
committed
update patching to base on OTel v1.33.x-0.54bx
1 parent 78568b2 commit a23151a

File tree

2 files changed

+106
-30
lines changed

2 files changed

+106
-30
lines changed

‎aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py

Lines changed: 36 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
_BotocoreInstrumentorContext,
4444
_BotoResultT,
4545
)
46+
from opentelemetry.instrumentation.botocore.utils import get_server_attributes
4647
from opentelemetry.instrumentation.utils import (
4748
is_instrumentation_enabled,
4849
suppress_http_instrumentation,
@@ -238,8 +239,8 @@ def _apply_botocore_dynamodb_patch() -> None:
238239
"""
239240
old_on_success = _DynamoDbExtension.on_success
240241

241-
def patch_on_success(self, span: Span, result: _BotoResultT):
242-
old_on_success(self, span, result)
242+
def patch_on_success(self, span: Span, result: _BotoResultT, instrumentor_context: _BotocoreInstrumentorContext):
243+
old_on_success(self, span, result, instrumentor_context)
243244
table = result.get("Table", {})
244245
table_arn = table.get("TableArn")
245246
if table_arn:
@@ -250,6 +251,19 @@ def patch_on_success(self, span: Span, result: _BotoResultT):
250251

251252
def _apply_botocore_api_call_patch() -> None:
252253
def patched_api_call(self, original_func, instance, args, kwargs):
254+
"""Botocore instrumentation patch to capture AWS authentication details
255+
256+
This patch extends the upstream implementation to include additional AWS authentication
257+
attributes:
258+
- aws.auth.account.access_key
259+
- aws.auth.region
260+
261+
Note: Current implementation duplicates upstream code in v1.33.x-0.54bx. Future improvements should:
262+
1. Propose refactoring upstream _patched_api_call into smaller components
263+
2. Apply targeted patches to these components to reduce code duplication
264+
265+
Reference: https://github.com/open-telemetry/opentelemetry-python-contrib/blob/release/v1.33.x-0.54bx/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py#L263
266+
"""
253267
if not is_instrumentation_enabled():
254268
return original_func(*args, **kwargs)
255269

@@ -265,7 +279,9 @@ def patched_api_call(self, original_func, instance, args, kwargs):
265279
SpanAttributes.RPC_SYSTEM: "aws-api",
266280
SpanAttributes.RPC_SERVICE: call_context.service_id,
267281
SpanAttributes.RPC_METHOD: call_context.operation,
282+
# TODO: update when semantic conventions exist
268283
"aws.region": call_context.region,
284+
**get_server_attributes(call_context.endpoint_url),
269285
AWS_AUTH_REGION: call_context.region,
270286
}
271287
credentials = instance._get_credentials()
@@ -276,13 +292,25 @@ def patched_api_call(self, original_func, instance, args, kwargs):
276292
attributes[AWS_AUTH_ACCESS_KEY] = access_key
277293

278294
_safe_invoke(extension.extract_attributes, attributes)
279-
280-
with self._tracer.start_as_current_span(
295+
end_span_on_exit = extension.should_end_span_on_exit()
296+
297+
tracer = self._get_tracer(extension)
298+
event_logger = self._get_event_logger(extension)
299+
meter = self._get_meter(extension)
300+
metrics = self._get_metrics(extension, meter)
301+
instrumentor_ctx = _BotocoreInstrumentorContext(
302+
event_logger=event_logger,
303+
metrics=metrics,
304+
)
305+
with tracer.start_as_current_span(
281306
call_context.span_name,
282307
kind=call_context.span_kind,
283308
attributes=attributes,
309+
# tracing streaming services require to close the span manually
310+
# at a later time after the stream has been consumed
311+
end_on_exit=end_span_on_exit,
284312
) as span:
285-
_safe_invoke(extension.before_service_call, span)
313+
_safe_invoke(extension.before_service_call, span, instrumentor_ctx)
286314
self._call_request_hook(span, call_context)
287315

288316
try:
@@ -293,12 +321,12 @@ def patched_api_call(self, original_func, instance, args, kwargs):
293321
except ClientError as error:
294322
result = getattr(error, "response", None)
295323
_apply_response_attributes(span, result)
296-
_safe_invoke(extension.on_error, span, error)
324+
_safe_invoke(extension.on_error, span, error, instrumentor_ctx)
297325
raise
298326
_apply_response_attributes(span, result)
299-
_safe_invoke(extension.on_success, span, result)
327+
_safe_invoke(extension.on_success, span, result, instrumentor_ctx)
300328
finally:
301-
_safe_invoke(extension.after_service_call)
329+
_safe_invoke(extension.after_service_call, instrumentor_ctx)
302330
self._call_response_hook(span, call_context, result)
303331

304332
return result

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

Lines changed: 70 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
AWS_BEDROCK_AGENT_ID,
1515
AWS_BEDROCK_DATA_SOURCE_ID,
1616
AWS_BEDROCK_GUARDRAIL_ARN,
17-
AWS_BEDROCK_GUARDRAIL_ARN,
1817
AWS_BEDROCK_GUARDRAIL_ID,
1918
AWS_BEDROCK_KNOWLEDGE_BASE_ID,
2019
AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
@@ -1034,14 +1033,24 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
10341033
# Validate behaviour of aws bucket name attribute, then remove it.
10351034
self._mock_attribute([SpanAttributes.AWS_S3_BUCKET], ["aws_s3_bucket_name"], keys, values)
10361035
self._validate_remote_resource_attributes(
1037-
"AWS::S3::Bucket", "aws_s3_bucket_name", None, _AWS_REMOTE_RESOURCE_REGION, None, _AWS_REMOTE_RESOURCE_ACCESS_KEY
1036+
"AWS::S3::Bucket",
1037+
"aws_s3_bucket_name",
1038+
None,
1039+
_AWS_REMOTE_RESOURCE_REGION,
1040+
None,
1041+
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
10381042
)
10391043
self._mock_attribute([SpanAttributes.AWS_S3_BUCKET], [None])
10401044

10411045
# Validate behaviour of AWS_SQS_QUEUE_NAME attribute, then remove it
10421046
self._mock_attribute([AWS_SQS_QUEUE_NAME], ["aws_queue_name"], keys, values)
10431047
self._validate_remote_resource_attributes(
1044-
"AWS::SQS::Queue", "aws_queue_name", None, _AWS_REMOTE_RESOURCE_REGION, None, _AWS_REMOTE_RESOURCE_ACCESS_KEY
1048+
"AWS::SQS::Queue",
1049+
"aws_queue_name",
1050+
None,
1051+
_AWS_REMOTE_RESOURCE_REGION,
1052+
None,
1053+
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
10451054
)
10461055
self._mock_attribute([AWS_SQS_QUEUE_NAME], [None])
10471056

@@ -1082,7 +1091,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
10821091
self._validate_remote_resource_attributes(
10831092
"AWS::Kinesis::Stream",
10841093
"aws_stream_name",
1085-
None,
1094+
None,
10861095
"us-west-2",
10871096
"123456789012",
10881097
None,
@@ -1094,7 +1103,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
10941103
self._validate_remote_resource_attributes(
10951104
"AWS::Kinesis::Stream",
10961105
"aws_stream_name",
1097-
None,
1106+
None,
10981107
_AWS_REMOTE_RESOURCE_REGION,
10991108
None,
11001109
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
@@ -1106,14 +1115,19 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
11061115
[AWS_DYNAMODB_TABLE_ARN], ["arn:aws:dynamodb:us-west-2:123456789012:table/aws_table_name"], keys, values
11071116
)
11081117
self._validate_remote_resource_attributes(
1109-
"AWS::DynamoDB::Table", "aws_table_name", None, "us-west-2", "123456789012", None
1118+
"AWS::DynamoDB::Table", "aws_table_name", None, "us-west-2", "123456789012", None
11101119
)
11111120
self._mock_attribute([AWS_DYNAMODB_TABLE_ARN], [None])
11121121

11131122
# Validate behaviour of SpanAttributes.AWS_DYNAMODB_TABLE_NAMES attribute with one table name, then remove it.
11141123
self._mock_attribute([SpanAttributes.AWS_DYNAMODB_TABLE_NAMES], [["aws_table_name"]], keys, values)
11151124
self._validate_remote_resource_attributes(
1116-
"AWS::DynamoDB::Table", "aws_table_name", None, _AWS_REMOTE_RESOURCE_REGION, None, _AWS_REMOTE_RESOURCE_ACCESS_KEY
1125+
"AWS::DynamoDB::Table",
1126+
"aws_table_name",
1127+
None,
1128+
_AWS_REMOTE_RESOURCE_REGION,
1129+
None,
1130+
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
11171131
)
11181132
self._mock_attribute([SpanAttributes.AWS_DYNAMODB_TABLE_NAMES], [None])
11191133

@@ -1134,7 +1148,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
11341148
self._validate_remote_resource_attributes(
11351149
"AWS::DynamoDB::Table",
11361150
"aws_table^|name",
1137-
None,
1151+
None,
11381152
_AWS_REMOTE_RESOURCE_REGION,
11391153
None,
11401154
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
@@ -1146,7 +1160,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
11461160
self._validate_remote_resource_attributes(
11471161
"AWS::DynamoDB::Table",
11481162
"aws_table^^name",
1149-
None,
1163+
None,
11501164
_AWS_REMOTE_RESOURCE_REGION,
11511165
None,
11521166
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
@@ -1156,14 +1170,24 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
11561170
# Validate behaviour of AWS_BEDROCK_AGENT_ID attribute, then remove it.
11571171
self._mock_attribute([AWS_BEDROCK_AGENT_ID], ["test_agent_id"], keys, values)
11581172
self._validate_remote_resource_attributes(
1159-
"AWS::Bedrock::Agent", "test_agent_id", None, _AWS_REMOTE_RESOURCE_REGION, None, _AWS_REMOTE_RESOURCE_ACCESS_KEY
1173+
"AWS::Bedrock::Agent",
1174+
"test_agent_id",
1175+
None,
1176+
_AWS_REMOTE_RESOURCE_REGION,
1177+
None,
1178+
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
11601179
)
11611180
self._mock_attribute([AWS_BEDROCK_AGENT_ID], [None])
11621181

11631182
# Validate behaviour of AWS_BEDROCK_AGENT_ID attribute with special chars(^), then remove it.
11641183
self._mock_attribute([AWS_BEDROCK_AGENT_ID], ["test_agent_^id"], keys, values)
11651184
self._validate_remote_resource_attributes(
1166-
"AWS::Bedrock::Agent", "test_agent_^^id", None, _AWS_REMOTE_RESOURCE_REGION, None, _AWS_REMOTE_RESOURCE_ACCESS_KEY
1185+
"AWS::Bedrock::Agent",
1186+
"test_agent_^^id",
1187+
None,
1188+
_AWS_REMOTE_RESOURCE_REGION,
1189+
None,
1190+
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
11671191
)
11681192
self._mock_attribute([AWS_BEDROCK_AGENT_ID], [None])
11691193

@@ -1240,7 +1264,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
12401264
self._validate_remote_resource_attributes(
12411265
"AWS::Bedrock::KnowledgeBase",
12421266
"test_knowledgeBase_id",
1243-
None,
1267+
None,
12441268
_AWS_REMOTE_RESOURCE_REGION,
12451269
None,
12461270
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
@@ -1252,7 +1276,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
12521276
self._validate_remote_resource_attributes(
12531277
"AWS::Bedrock::KnowledgeBase",
12541278
"test_knowledgeBase_^^id",
1255-
None,
1279+
None,
12561280
_AWS_REMOTE_RESOURCE_REGION,
12571281
None,
12581282
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
@@ -1262,7 +1286,12 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
12621286
# Validate behaviour of GEN_AI_REQUEST_MODEL attribute, then remove it.
12631287
self._mock_attribute([GEN_AI_REQUEST_MODEL], ["test.service_id"], keys, values)
12641288
self._validate_remote_resource_attributes(
1265-
"AWS::Bedrock::Model", "test.service_id", None, _AWS_REMOTE_RESOURCE_REGION, None, _AWS_REMOTE_RESOURCE_ACCESS_KEY
1289+
"AWS::Bedrock::Model",
1290+
"test.service_id",
1291+
None,
1292+
_AWS_REMOTE_RESOURCE_REGION,
1293+
None,
1294+
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
12661295
)
12671296
self._mock_attribute([GEN_AI_REQUEST_MODEL], [None])
12681297

@@ -1271,7 +1300,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
12711300
self._validate_remote_resource_attributes(
12721301
"AWS::Bedrock::Model",
12731302
"test.service_^^id",
1274-
None,
1303+
None,
12751304
_AWS_REMOTE_RESOURCE_REGION,
12761305
None,
12771306
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
@@ -1346,7 +1375,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
13461375
self._validate_remote_resource_attributes(
13471376
"AWS::Lambda::EventSourceMapping",
13481377
"aws_event_source_mapping_id",
1349-
None,
1378+
None,
13501379
_AWS_REMOTE_RESOURCE_REGION,
13511380
None,
13521381
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
@@ -1364,7 +1393,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
13641393
self._validate_remote_resource_attributes(
13651394
"AWS::Lambda::EventSourceMapping",
13661395
"aws_event_source_mapping_id",
1367-
None,
1396+
None,
13681397
_AWS_REMOTE_RESOURCE_REGION,
13691398
None,
13701399
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
@@ -1449,7 +1478,12 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
14491478
values,
14501479
)
14511480
self._validate_remote_resource_attributes(
1452-
"AWS::Lambda::Function", "testLambdaName", "arn:aws:lambda:us-east-1:123456789012:function:testLambdaName", "us-east-1", "123456789012", None
1481+
"AWS::Lambda::Function",
1482+
"testLambdaName",
1483+
"arn:aws:lambda:us-east-1:123456789012:function:testLambdaName",
1484+
"us-east-1",
1485+
"123456789012",
1486+
None,
14531487
)
14541488
self._mock_attribute(
14551489
[
@@ -1496,7 +1530,7 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
14961530
self._validate_remote_resource_attributes(
14971531
"AWS::StepFunctions::StateMachine",
14981532
"invalid_arn",
1499-
None,
1533+
None,
15001534
_AWS_REMOTE_RESOURCE_REGION,
15011535
None,
15021536
_AWS_REMOTE_RESOURCE_ACCESS_KEY,
@@ -1523,7 +1557,12 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
15231557
["arn:aws:states:us-east-1:123456789123:stateMachine:testStateMachine", "aws-api"],
15241558
)
15251559
self._validate_remote_resource_attributes(
1526-
"AWS::StepFunctions::StateMachine", "testStateMachine", "arn:aws:states:us-east-1:123456789123:stateMachine:testStateMachine", "us-east-1", "123456789123", None
1560+
"AWS::StepFunctions::StateMachine",
1561+
"testStateMachine",
1562+
"arn:aws:states:us-east-1:123456789123:stateMachine:testStateMachine",
1563+
"us-east-1",
1564+
"123456789123",
1565+
None,
15271566
)
15281567
self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None])
15291568

@@ -1532,7 +1571,11 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
15321571
[AWS_STEPFUNCTIONS_STATEMACHINE_ARN, SpanAttributes.RPC_SYSTEM],
15331572
["arn:aws:states:us-east-1:invalid_account_id:stateMachine:testStateMachine", "aws-api"],
15341573
)
1535-
self._validate_remote_resource_attributes("AWS::StepFunctions::StateMachine", "testStateMachine", "arn:aws:states:us-east-1:invalid_account_id:stateMachine:testStateMachine")
1574+
self._validate_remote_resource_attributes(
1575+
"AWS::StepFunctions::StateMachine",
1576+
"testStateMachine",
1577+
"arn:aws:states:us-east-1:invalid_account_id:stateMachine:testStateMachine",
1578+
)
15361579
self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None])
15371580

15381581
# Arn with invalid region
@@ -1541,7 +1584,12 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
15411584
["arn:aws:states:invalid_region:123456789123:stateMachine:testStateMachine", "aws-api"],
15421585
)
15431586
self._validate_remote_resource_attributes(
1544-
"AWS::StepFunctions::StateMachine", "testStateMachine", "arn:aws:states:invalid_region:123456789123:stateMachine:testStateMachine", "invalid_region", "123456789123", None
1587+
"AWS::StepFunctions::StateMachine",
1588+
"testStateMachine",
1589+
"arn:aws:states:invalid_region:123456789123:stateMachine:testStateMachine",
1590+
"invalid_region",
1591+
"123456789123",
1592+
None,
15451593
)
15461594
self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None])
15471595

0 commit comments

Comments
 (0)