Skip to content

Commit c690e5d

Browse files
pabloDeputterpabloDeputter
committed
feat(boto3): Add Kinesis extension
1 parent 009f15f commit c690e5d

4 files changed

Lines changed: 152 additions & 14 deletions

File tree

‎sentry_sdk/consts.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,12 @@ class SPANDATA:
363363
Example: "wzHcyEWfmOGDIE5QOhTAqFDoDWP3y8IUvpNINCwL9N4TEHbUw0/gZJ+VZTmCNCWR7fezEN3eCiQ="
364364
"""
365365

366+
AWS_KINESIS_STREAM_NAME = "aws.kinesis.stream_name"
367+
"""
368+
The name of the AWS Kinesis stream the request refers to.
369+
Example: "orders"
370+
"""
371+
366372
AWS_LAMBDA_INVOKED_ARN = "aws.lambda.invoked_arn"
367373
"""
368374
The full ARN of the Lambda function that was invoked.

‎sentry_sdk/integrations/boto3/_client.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from typing import TYPE_CHECKING
33

44
import sentry_sdk
5+
from sentry_sdk.consts import SPANDATA
56
from sentry_sdk.integrations import DidNotEnable
67
from sentry_sdk.integrations.boto3._context import AwsCallContext
78
from sentry_sdk.integrations.boto3._instrumentation import (
@@ -17,7 +18,7 @@
1718
_resolve_service,
1819
)
1920
from sentry_sdk.integrations.boto3.consts import IDENTIFIER
20-
from sentry_sdk.traces import NoOpSpan, Span
21+
from sentry_sdk.traces import NoOpSpan, Span, SpanStatus
2122
from sentry_sdk.utils import capture_internal_exceptions
2223

2324
if TYPE_CHECKING:
@@ -125,6 +126,8 @@ def sentry_patched_make_api_call(
125126
attributes.update(_get_response_attributes(parsed))
126127
with capture_internal_exceptions():
127128
span.set_attributes(attributes)
129+
if SPANDATA.ERROR_TYPE in attributes:
130+
span.status = SpanStatus.ERROR
128131
except BaseException as error:
129132
_finish_span(span, error)
130133
raise

‎sentry_sdk/integrations/boto3/_services/kinesis.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,30 @@
11
from typing import TYPE_CHECKING
22

3+
from sentry_sdk.consts import SPANDATA
34
from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
45
from sentry_sdk.integrations.boto3._utils import _get_aws_arn_attributes
56
from sentry_sdk.utils import capture_internal_exceptions
67

78
if TYPE_CHECKING:
9+
from typing import Any
10+
811
from sentry_sdk._types import Attributes
912
from sentry_sdk.integrations.boto3._context import AwsCallContext
1013

1114

1215
class _KinesisExtension(_ServiceExtension):
1316
def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
1417
attributes: "Attributes" = {}
18+
19+
# a stream name, or an arn that contains it: `…:stream/<name>` for a stream,
20+
# `…:stream/<name>/consumer/<consumer>:<timestamp>` for a consumer.
21+
for param in ("StreamName", "StreamARN", "ConsumerARN", "ResourceARN"):
22+
if param in ctx.params:
23+
attributes[SPANDATA.AWS_KINESIS_STREAM_NAME] = (
24+
ctx.params[param].split(":stream/", 1)[-1].split("/", 1)[0]
25+
)
26+
break
27+
1528
with capture_internal_exceptions():
1629
attributes.update(
1730
_get_aws_arn_attributes(
@@ -24,4 +37,24 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
2437
("kinesis",),
2538
)
2639
)
40+
2741
return attributes
42+
43+
def get_response_attributes(
44+
self, ctx: "AwsCallContext", response: "Any"
45+
) -> "Attributes":
46+
# `PutRecords` returns `200 OK` even when records fail.
47+
# https://docs.aws.amazon.com/kinesis/latest/APIReference/API_PutRecords.html
48+
if ctx.operation_name != "PutRecords" or not response.get("FailedRecordCount"):
49+
return {}
50+
51+
error_codes = {
52+
record["ErrorCode"]
53+
for record in response["Records"]
54+
if "ErrorCode" in record
55+
}
56+
return {
57+
SPANDATA.ERROR_TYPE: error_codes.pop()
58+
if len(error_codes) == 1
59+
else "aws_batch_partial_failure"
60+
}

‎tests/integrations/boto3/test_kinesis.py‎

Lines changed: 109 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,120 @@
1+
import pytest
12
from botocore.stub import Stubber
23

34
from sentry_sdk.consts import OP, SPANDATA
4-
from tests.integrations.boto3.helpers import capture_spans_by_op
5+
from tests.integrations.boto3.helpers import (
6+
capture_spans_by_op,
7+
require_botocore_model_fields,
8+
)
59
from tests.integrations.boto3.helpers import client_factory as client_factory
610

11+
STREAM_ARN = "arn:aws:kinesis:eu-north-1:123456789012:stream/orders"
12+
CONSUMER_ARN = STREAM_ARN + "/consumer/billing:1760000000"
713

8-
def test_request_attributes(client_factory, capture_items):
9-
client = client_factory("kinesis")
10-
stream_arn = "arn:aws:kinesis:eu-north-1:123456789012:stream/orders"
1114

15+
def _stubbed_span(client_factory, capture_items, method, params, response):
16+
client = client_factory("kinesis")
17+
require_botocore_model_fields(
18+
client, method, input_fields=tuple(params), output_fields=tuple(response)
19+
)
1220
with Stubber(client) as stubber:
13-
stubber.add_response(
14-
"list_stream_consumers", {"Consumers": []}, {"StreamARN": stream_arn}
15-
)
21+
stubber.add_response(method, response, params)
1622
spans = capture_spans_by_op(
17-
lambda: client.list_stream_consumers(**{"StreamARN": stream_arn}),
18-
capture_items,
23+
lambda: getattr(client, method)(**params), capture_items
1924
)
20-
2125
(span,) = spans[OP.RPC]
22-
attributes = span["attributes"]
23-
assert attributes[SPANDATA.CLOUD_ACCOUNT_ID] == "123456789012"
24-
assert attributes[SPANDATA.CLOUD_RESOURCE_ID] == stream_arn
26+
return span
27+
28+
29+
@pytest.mark.parametrize(
30+
"method,params,response,expected",
31+
[
32+
pytest.param(
33+
"put_record",
34+
{"StreamName": "orders", "Data": b"{}", "PartitionKey": "customer-1"},
35+
{"ShardId": "shardId-000000000000", "SequenceNumber": "1"},
36+
{SPANDATA.AWS_KINESIS_STREAM_NAME: "orders"},
37+
id="stream-name",
38+
),
39+
pytest.param(
40+
"list_stream_consumers",
41+
{"StreamARN": STREAM_ARN},
42+
{"Consumers": []},
43+
{
44+
SPANDATA.AWS_KINESIS_STREAM_NAME: "orders",
45+
SPANDATA.CLOUD_RESOURCE_ID: STREAM_ARN,
46+
},
47+
id="stream-arn",
48+
),
49+
pytest.param(
50+
"deregister_stream_consumer",
51+
{"ConsumerARN": CONSUMER_ARN},
52+
{},
53+
{
54+
SPANDATA.AWS_KINESIS_STREAM_NAME: "orders",
55+
SPANDATA.CLOUD_RESOURCE_ID: CONSUMER_ARN,
56+
},
57+
id="consumer-arn",
58+
),
59+
pytest.param(
60+
"list_tags_for_resource",
61+
{"ResourceARN": CONSUMER_ARN},
62+
{"Tags": []},
63+
{
64+
SPANDATA.AWS_KINESIS_STREAM_NAME: "orders",
65+
SPANDATA.CLOUD_RESOURCE_ID: CONSUMER_ARN,
66+
},
67+
id="resource-arn",
68+
),
69+
],
70+
)
71+
def test_stream_attributes(
72+
client_factory, capture_items, method, params, response, expected
73+
):
74+
span = _stubbed_span(client_factory, capture_items, method, params, response)
75+
76+
assert {
77+
key: value
78+
for key, value in span["attributes"].items()
79+
if key in (SPANDATA.AWS_KINESIS_STREAM_NAME, SPANDATA.CLOUD_RESOURCE_ID)
80+
} == expected
81+
82+
83+
@pytest.mark.parametrize(
84+
"error_codes,expected_error_type",
85+
[
86+
pytest.param((), None, id="ok"),
87+
pytest.param(
88+
("ProvisionedThroughputExceededException",) * 2,
89+
"ProvisionedThroughputExceededException",
90+
id="one-code",
91+
),
92+
pytest.param(
93+
("ProvisionedThroughputExceededException", "InternalFailure"),
94+
"aws_batch_partial_failure",
95+
id="different-codes",
96+
),
97+
],
98+
)
99+
def test_put_records_outcome(
100+
client_factory, capture_items, error_codes, expected_error_type
101+
):
102+
params = {
103+
"StreamName": "orders",
104+
"Records": [
105+
{"Data": b"card 4242", "PartitionKey": "customer-1"},
106+
{"Data": b"card 4242", "PartitionKey": "customer-2"},
107+
{"Data": b"card 4242", "PartitionKey": "customer-3"},
108+
],
109+
}
110+
response = {
111+
"Records": [{"ShardId": "shardId-000000000000", "SequenceNumber": "1"}]
112+
+ [{"ErrorCode": code, "ErrorMessage": "failed"} for code in error_codes],
113+
}
114+
# the model's minimum is 1, so the stubber only accepts it for failed records.
115+
if error_codes:
116+
response["FailedRecordCount"] = len(error_codes) # type: ignore
117+
span = _stubbed_span(client_factory, capture_items, "put_records", params, response)
118+
119+
assert span["status"] == ("ok" if expected_error_type is None else "error")
120+
assert span["attributes"].get(SPANDATA.ERROR_TYPE) == expected_error_type

0 commit comments

Comments
 (0)