Skip to content

Commit ba1e550

Browse files
marshall-leemarshall-lee
authored
Boto3 integration (getsentry#896)
This is the integration for boto3 library for recording AWS requests as spans. Another suggestion is to enable it by default in aws_lambda integration since boto3 package is pre-installed on every lambda.
1 parent 881b8e1 commit ba1e550

7 files changed

Lines changed: 259 additions & 0 deletions

File tree

‎sentry_sdk/integrations/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ def iter_default_integrations(with_auto_enabling_integrations):
6262
"sentry_sdk.integrations.aiohttp.AioHttpIntegration",
6363
"sentry_sdk.integrations.tornado.TornadoIntegration",
6464
"sentry_sdk.integrations.sqlalchemy.SqlalchemyIntegration",
65+
"sentry_sdk.integrations.boto3.Boto3Integration",
6566
)
6667

6768

‎sentry_sdk/integrations/boto3.py‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
from __future__ import absolute_import
2+
3+
from sentry_sdk import Hub
4+
from sentry_sdk.integrations import Integration, DidNotEnable
5+
from sentry_sdk.tracing import Span
6+
7+
from sentry_sdk._functools import partial
8+
from sentry_sdk._types import MYPY
9+
10+
if MYPY:
11+
from typing import Any
12+
from typing import Dict
13+
from typing import Optional
14+
from typing import Type
15+
16+
try:
17+
from botocore.client import BaseClient # type: ignore
18+
from botocore.response import StreamingBody # type: ignore
19+
from botocore.awsrequest import AWSRequest # type: ignore
20+
except ImportError:
21+
raise DidNotEnable("botocore is not installed")
22+
23+
24+
class Boto3Integration(Integration):
25+
identifier = "boto3"
26+
27+
@staticmethod
28+
def setup_once():
29+
# type: () -> None
30+
orig_init = BaseClient.__init__
31+
32+
def sentry_patched_init(self, *args, **kwargs):
33+
# type: (Type[BaseClient], *Any, **Any) -> None
34+
orig_init(self, *args, **kwargs)
35+
meta = self.meta
36+
service_id = meta.service_model.service_id.hyphenize()
37+
meta.events.register(
38+
"request-created",
39+
partial(_sentry_request_created, service_id=service_id),
40+
)
41+
meta.events.register("after-call", _sentry_after_call)
42+
meta.events.register("after-call-error", _sentry_after_call_error)
43+
44+
BaseClient.__init__ = sentry_patched_init
45+
46+
47+
def _sentry_request_created(service_id, request, operation_name, **kwargs):
48+
# type: (str, AWSRequest, str, **Any) -> None
49+
hub = Hub.current
50+
if hub.get_integration(Boto3Integration) is None:
51+
return
52+
53+
description = "aws.%s.%s" % (service_id, operation_name)
54+
span = hub.start_span(
55+
hub=hub,
56+
op="aws.request",
57+
description=description,
58+
)
59+
span.set_tag("aws.service_id", service_id)
60+
span.set_tag("aws.operation_name", operation_name)
61+
span.set_data("aws.request.url", request.url)
62+
63+
# We do it in order for subsequent http calls/retries be
64+
# attached to this span.
65+
span.__enter__()
66+
67+
# request.context is an open-ended data-structure
68+
# where we can add anything useful in request life cycle.
69+
request.context["_sentrysdk_span"] = span
70+
71+
72+
def _sentry_after_call(context, parsed, **kwargs):
73+
# type: (Dict[str, Any], Dict[str, Any], **Any) -> None
74+
span = context.pop("_sentrysdk_span", None) # type: Optional[Span]
75+
76+
# Span could be absent if the integration is disabled.
77+
if span is None:
78+
return
79+
span.__exit__(None, None, None)
80+
81+
body = parsed.get("Body")
82+
if not isinstance(body, StreamingBody):
83+
return
84+
85+
streaming_span = span.start_child(
86+
op="aws.request.stream",
87+
description=span.description,
88+
)
89+
90+
orig_read = body.read
91+
orig_close = body.close
92+
93+
def sentry_streaming_body_read(*args, **kwargs):
94+
# type: (*Any, **Any) -> bytes
95+
try:
96+
ret = orig_read(*args, **kwargs)
97+
if not ret:
98+
streaming_span.finish()
99+
return ret
100+
except Exception:
101+
streaming_span.finish()
102+
raise
103+
104+
body.read = sentry_streaming_body_read
105+
106+
def sentry_streaming_body_close(*args, **kwargs):
107+
# type: (*Any, **Any) -> None
108+
streaming_span.finish()
109+
orig_close(*args, **kwargs)
110+
111+
body.close = sentry_streaming_body_close
112+
113+
114+
def _sentry_after_call_error(context, exception, **kwargs):
115+
# type: (Dict[str, Any], Type[BaseException], **Any) -> None
116+
span = context.pop("_sentrysdk_span", None) # type: Optional[Span]
117+
118+
# Span could be absent if the integration is disabled.
119+
if span is None:
120+
return
121+
span.__exit__(type(exception), exception, None)

‎tests/integrations/boto3/__init__.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
import pytest
2+
import os
3+
4+
pytest.importorskip("boto3")
5+
xml_fixture_path = os.path.dirname(os.path.abspath(__file__))
6+
7+
8+
def read_fixture(name):
9+
with open(os.path.join(xml_fixture_path, name), "rb") as f:
10+
return f.read()

‎tests/integrations/boto3/aws_mock.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
from io import BytesIO
2+
from botocore.awsrequest import AWSResponse
3+
4+
5+
class Body(BytesIO):
6+
def stream(self, **kwargs):
7+
contents = self.read()
8+
while contents:
9+
yield contents
10+
contents = self.read()
11+
12+
13+
class MockResponse(object):
14+
def __init__(self, client, status_code, headers, body):
15+
self._client = client
16+
self._status_code = status_code
17+
self._headers = headers
18+
self._body = body
19+
20+
def __enter__(self):
21+
self._client.meta.events.register("before-send", self)
22+
return self
23+
24+
def __exit__(self, exc_type, exc_value, traceback):
25+
self._client.meta.events.unregister("before-send", self)
26+
27+
def __call__(self, request, **kwargs):
28+
return AWSResponse(
29+
request.url,
30+
self._status_code,
31+
self._headers,
32+
Body(self._body),
33+
)

‎tests/integrations/boto3/s3_list.xml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/"><Name>marshalls-furious-bucket</Name><Prefix></Prefix><Marker></Marker><MaxKeys>1000</MaxKeys><EncodingType>url</EncodingType><IsTruncated>false</IsTruncated><Contents><Key>foo.txt</Key><LastModified>2020-10-24T00:13:39.000Z</LastModified><ETag>&quot;a895ba674b4abd01b5d67cfd7074b827&quot;</ETag><Size>206453</Size><Owner><ID>7bef397f7e536914d1ff1bbdb105ed90bcfd06269456bf4a06c6e2e54564daf7</ID></Owner><StorageClass>STANDARD</StorageClass></Contents><Contents><Key>bar.txt</Key><LastModified>2020-10-02T15:15:20.000Z</LastModified><ETag>&quot;a895ba674b4abd01b5d67cfd7074b827&quot;</ETag><Size>206453</Size><Owner><ID>7bef397f7e536914d1ff1bbdb105ed90bcfd06269456bf4a06c6e2e54564daf7</ID></Owner><StorageClass>STANDARD</StorageClass></Contents></ListBucketResult>

‎tests/integrations/boto3/test_s3.py‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
from sentry_sdk import Hub
2+
from sentry_sdk.integrations.boto3 import Boto3Integration
3+
from tests.integrations.boto3.aws_mock import MockResponse
4+
from tests.integrations.boto3 import read_fixture
5+
6+
import boto3
7+
8+
session = boto3.Session(
9+
aws_access_key_id="-",
10+
aws_secret_access_key="-",
11+
)
12+
13+
14+
def test_basic(sentry_init, capture_events):
15+
sentry_init(traces_sample_rate=1.0, integrations=[Boto3Integration()])
16+
events = capture_events()
17+
18+
s3 = session.resource("s3")
19+
with Hub.current.start_transaction() as transaction, MockResponse(
20+
s3.meta.client, 200, {}, read_fixture("s3_list.xml")
21+
):
22+
bucket = s3.Bucket("bucket")
23+
items = [obj for obj in bucket.objects.all()]
24+
assert len(items) == 2
25+
assert items[0].key == "foo.txt"
26+
assert items[1].key == "bar.txt"
27+
transaction.finish()
28+
29+
(event,) = events
30+
assert event["type"] == "transaction"
31+
assert len(event["spans"]) == 1
32+
(span,) = event["spans"]
33+
assert span["op"] == "aws.request"
34+
assert span["description"] == "aws.s3.ListObjects"
35+
36+
37+
def test_streaming(sentry_init, capture_events):
38+
sentry_init(traces_sample_rate=1.0, integrations=[Boto3Integration()])
39+
events = capture_events()
40+
41+
s3 = session.resource("s3")
42+
with Hub.current.start_transaction() as transaction, MockResponse(
43+
s3.meta.client, 200, {}, b"hello"
44+
):
45+
obj = s3.Bucket("bucket").Object("foo.pdf")
46+
body = obj.get()["Body"]
47+
assert body.read(1) == b"h"
48+
assert body.read(2) == b"el"
49+
assert body.read(3) == b"lo"
50+
assert body.read(1) == b""
51+
transaction.finish()
52+
53+
(event,) = events
54+
assert event["type"] == "transaction"
55+
assert len(event["spans"]) == 2
56+
span1 = event["spans"][0]
57+
assert span1["op"] == "aws.request"
58+
assert span1["description"] == "aws.s3.GetObject"
59+
span2 = event["spans"][1]
60+
assert span2["op"] == "aws.request.stream"
61+
assert span2["description"] == "aws.s3.GetObject"
62+
assert span2["parent_span_id"] == span1["span_id"]
63+
64+
65+
def test_streaming_close(sentry_init, capture_events):
66+
sentry_init(traces_sample_rate=1.0, integrations=[Boto3Integration()])
67+
events = capture_events()
68+
69+
s3 = session.resource("s3")
70+
with Hub.current.start_transaction() as transaction, MockResponse(
71+
s3.meta.client, 200, {}, b"hello"
72+
):
73+
obj = s3.Bucket("bucket").Object("foo.pdf")
74+
body = obj.get()["Body"]
75+
assert body.read(1) == b"h"
76+
body.close() # close partially-read stream
77+
transaction.finish()
78+
79+
(event,) = events
80+
assert event["type"] == "transaction"
81+
assert len(event["spans"]) == 2
82+
span1 = event["spans"][0]
83+
assert span1["op"] == "aws.request"
84+
span2 = event["spans"][1]
85+
assert span2["op"] == "aws.request.stream"

‎tox.ini‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,8 @@ envlist =
8181

8282
{py3.6,py3.7,py3.8}-chalice-{1.16,1.17,1.18,1.19,1.20}
8383

84+
{py2.7,py3.6,py3.7,py3.8}-boto3-{1.14,1.15,1.16}
85+
8486
[testenv]
8587
deps =
8688
# if you change test-requirements.txt and your change is not being reflected
@@ -224,6 +226,10 @@ deps =
224226
chalice-1.20: chalice>=1.20.0,<1.21.0
225227
chalice: pytest-chalice==0.0.5
226228

229+
boto3-1.14: boto3>=1.14,<1.15
230+
boto3-1.15: boto3>=1.15,<1.16
231+
boto3-1.16: boto3>=1.16,<1.17
232+
227233
setenv =
228234
PYTHONDONTWRITEBYTECODE=1
229235
TESTPATH=tests
@@ -249,6 +255,7 @@ setenv =
249255
spark: TESTPATH=tests/integrations/spark
250256
pure_eval: TESTPATH=tests/integrations/pure_eval
251257
chalice: TESTPATH=tests/integrations/chalice
258+
boto3: TESTPATH=tests/integrations/boto3
252259

253260
COVERAGE_FILE=.coverage-{envname}
254261
passenv =

0 commit comments

Comments
 (0)