Boto3 integration (#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.

Author: marshall-lee

sentry_sdk/integrations/__init__.py

@@ -62,6 +62,7 @@ def iter_default_integrations(with_auto_enabling_integrations):
     "sentry_sdk.integrations.aiohttp.AioHttpIntegration",
     "sentry_sdk.integrations.tornado.TornadoIntegration",
     "sentry_sdk.integrations.sqlalchemy.SqlalchemyIntegration",
+    "sentry_sdk.integrations.boto3.Boto3Integration",
 )
 
 

sentry_sdk/integrations/boto3.py

@@ -0,0 +1,121 @@
+from __future__ import absolute_import
+
+from sentry_sdk import Hub
+from sentry_sdk.integrations import Integration, DidNotEnable
+from sentry_sdk.tracing import Span
+
+from sentry_sdk._functools import partial
+from sentry_sdk._types import MYPY
+
+if MYPY:
+    from typing import Any
+    from typing import Dict
+    from typing import Optional
+    from typing import Type
+
+try:
+    from botocore.client import BaseClient  # type: ignore
+    from botocore.response import StreamingBody  # type: ignore
+    from botocore.awsrequest import AWSRequest  # type: ignore
+except ImportError:
+    raise DidNotEnable("botocore is not installed")
+
+
+class Boto3Integration(Integration):
+    identifier = "boto3"
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        orig_init = BaseClient.__init__
+
+        def sentry_patched_init(self, *args, **kwargs):
+            # type: (Type[BaseClient], *Any, **Any) -> None
+            orig_init(self, *args, **kwargs)
+            meta = self.meta
+            service_id = meta.service_model.service_id.hyphenize()
+            meta.events.register(
+                "request-created",
+                partial(_sentry_request_created, service_id=service_id),
+            )
+            meta.events.register("after-call", _sentry_after_call)
+            meta.events.register("after-call-error", _sentry_after_call_error)
+
+        BaseClient.__init__ = sentry_patched_init
+
+
+def _sentry_request_created(service_id, request, operation_name, **kwargs):
+    # type: (str, AWSRequest, str, **Any) -> None
+    hub = Hub.current
+    if hub.get_integration(Boto3Integration) is None:
+        return
+
+    description = "aws.%s.%s" % (service_id, operation_name)
+    span = hub.start_span(
+        hub=hub,
+        op="aws.request",
+        description=description,
+    )
+    span.set_tag("aws.service_id", service_id)
+    span.set_tag("aws.operation_name", operation_name)
+    span.set_data("aws.request.url", request.url)
+
+    # We do it in order for subsequent http calls/retries be
+    # attached to this span.
+    span.__enter__()
+
+    # request.context is an open-ended data-structure
+    # where we can add anything useful in request life cycle.
+    request.context["_sentrysdk_span"] = span
+
+
+def _sentry_after_call(context, parsed, **kwargs):
+    # type: (Dict[str, Any], Dict[str, Any], **Any) -> None
+    span = context.pop("_sentrysdk_span", None)  # type: Optional[Span]
+
+    # Span could be absent if the integration is disabled.
+    if span is None:
+        return
+    span.__exit__(None, None, None)
+
+    body = parsed.get("Body")
+    if not isinstance(body, StreamingBody):
+        return
+
+    streaming_span = span.start_child(
+        op="aws.request.stream",
+        description=span.description,
+    )
+
+    orig_read = body.read
+    orig_close = body.close
+
+    def sentry_streaming_body_read(*args, **kwargs):
+        # type: (*Any, **Any) -> bytes
+        try:
+            ret = orig_read(*args, **kwargs)
+            if not ret:
+                streaming_span.finish()
+            return ret
+        except Exception:
+            streaming_span.finish()
+            raise
+
+    body.read = sentry_streaming_body_read
+
+    def sentry_streaming_body_close(*args, **kwargs):
+        # type: (*Any, **Any) -> None
+        streaming_span.finish()
+        orig_close(*args, **kwargs)
+
+    body.close = sentry_streaming_body_close
+
+
+def _sentry_after_call_error(context, exception, **kwargs):
+    # type: (Dict[str, Any], Type[BaseException], **Any) -> None
+    span = context.pop("_sentrysdk_span", None)  # type: Optional[Span]
+
+    # Span could be absent if the integration is disabled.
+    if span is None:
+        return
+    span.__exit__(type(exception), exception, None)

tests/integrations/boto3/__init__.py

@@ -0,0 +1,10 @@
+import pytest
+import os
+
+pytest.importorskip("boto3")
+xml_fixture_path = os.path.dirname(os.path.abspath(__file__))
+
+
+def read_fixture(name):
+    with open(os.path.join(xml_fixture_path, name), "rb") as f:
+        return f.read()

tests/integrations/boto3/aws_mock.py

@@ -0,0 +1,33 @@
+from io import BytesIO
+from botocore.awsrequest import AWSResponse
+
+
+class Body(BytesIO):
+    def stream(self, **kwargs):
+        contents = self.read()
+        while contents:
+            yield contents
+            contents = self.read()
+
+
+class MockResponse(object):
+    def __init__(self, client, status_code, headers, body):
+        self._client = client
+        self._status_code = status_code
+        self._headers = headers
+        self._body = body
+
+    def __enter__(self):
+        self._client.meta.events.register("before-send", self)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._client.meta.events.unregister("before-send", self)
+
+    def __call__(self, request, **kwargs):
+        return AWSResponse(
+            request.url,
+            self._status_code,
+            self._headers,
+            Body(self._body),
+        )

tests/integrations/boto3/s3_list.xml

@@ -0,0 +1,2 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<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

@@ -0,0 +1,85 @@
+from sentry_sdk import Hub
+from sentry_sdk.integrations.boto3 import Boto3Integration
+from tests.integrations.boto3.aws_mock import MockResponse
+from tests.integrations.boto3 import read_fixture
+
+import boto3
+
+session = boto3.Session(
+    aws_access_key_id="-",
+    aws_secret_access_key="-",
+)
+
+
+def test_basic(sentry_init, capture_events):
+    sentry_init(traces_sample_rate=1.0, integrations=[Boto3Integration()])
+    events = capture_events()
+
+    s3 = session.resource("s3")
+    with Hub.current.start_transaction() as transaction, MockResponse(
+        s3.meta.client, 200, {}, read_fixture("s3_list.xml")
+    ):
+        bucket = s3.Bucket("bucket")
+        items = [obj for obj in bucket.objects.all()]
+        assert len(items) == 2
+        assert items[0].key == "foo.txt"
+        assert items[1].key == "bar.txt"
+        transaction.finish()
+
+    (event,) = events
+    assert event["type"] == "transaction"
+    assert len(event["spans"]) == 1
+    (span,) = event["spans"]
+    assert span["op"] == "aws.request"
+    assert span["description"] == "aws.s3.ListObjects"
+
+
+def test_streaming(sentry_init, capture_events):
+    sentry_init(traces_sample_rate=1.0, integrations=[Boto3Integration()])
+    events = capture_events()
+
+    s3 = session.resource("s3")
+    with Hub.current.start_transaction() as transaction, MockResponse(
+        s3.meta.client, 200, {}, b"hello"
+    ):
+        obj = s3.Bucket("bucket").Object("foo.pdf")
+        body = obj.get()["Body"]
+        assert body.read(1) == b"h"
+        assert body.read(2) == b"el"
+        assert body.read(3) == b"lo"
+        assert body.read(1) == b""
+        transaction.finish()
+
+    (event,) = events
+    assert event["type"] == "transaction"
+    assert len(event["spans"]) == 2
+    span1 = event["spans"][0]
+    assert span1["op"] == "aws.request"
+    assert span1["description"] == "aws.s3.GetObject"
+    span2 = event["spans"][1]
+    assert span2["op"] == "aws.request.stream"
+    assert span2["description"] == "aws.s3.GetObject"
+    assert span2["parent_span_id"] == span1["span_id"]
+
+
+def test_streaming_close(sentry_init, capture_events):
+    sentry_init(traces_sample_rate=1.0, integrations=[Boto3Integration()])
+    events = capture_events()
+
+    s3 = session.resource("s3")
+    with Hub.current.start_transaction() as transaction, MockResponse(
+        s3.meta.client, 200, {}, b"hello"
+    ):
+        obj = s3.Bucket("bucket").Object("foo.pdf")
+        body = obj.get()["Body"]
+        assert body.read(1) == b"h"
+        body.close()  # close partially-read stream
+        transaction.finish()
+
+    (event,) = events
+    assert event["type"] == "transaction"
+    assert len(event["spans"]) == 2
+    span1 = event["spans"][0]
+    assert span1["op"] == "aws.request"
+    span2 = event["spans"][1]
+    assert span2["op"] == "aws.request.stream"

tox.ini

@@ -81,6 +81,8 @@ envlist =
 
     {py3.6,py3.7,py3.8}-chalice-{1.16,1.17,1.18,1.19,1.20}
 
+    {py2.7,py3.6,py3.7,py3.8}-boto3-{1.14,1.15,1.16}
+
 [testenv]
 deps =
     # if you change test-requirements.txt and your change is not being reflected
@@ -224,6 +226,10 @@ deps =
     chalice-1.20: chalice>=1.20.0,<1.21.0
     chalice: pytest-chalice==0.0.5
 
+    boto3-1.14: boto3>=1.14,<1.15
+    boto3-1.15: boto3>=1.15,<1.16
+    boto3-1.16: boto3>=1.16,<1.17
+
 setenv =
     PYTHONDONTWRITEBYTECODE=1
     TESTPATH=tests
@@ -249,6 +255,7 @@ setenv =
     spark: TESTPATH=tests/integrations/spark
     pure_eval: TESTPATH=tests/integrations/pure_eval
     chalice: TESTPATH=tests/integrations/chalice
+    boto3: TESTPATH=tests/integrations/boto3
 
     COVERAGE_FILE=.coverage-{envname}
 passenv =