From ba1e55009822a8dc8e231158254ea207bf3a5bab Mon Sep 17 00:00:00 2001 From: Vladimir Kochnev Date: Thu, 29 Oct 2020 15:35:10 +0000 Subject: [PATCH] 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. --- sentry_sdk/integrations/__init__.py | 1 + sentry_sdk/integrations/boto3.py | 121 +++++++++++++++++++++++++++ tests/integrations/boto3/__init__.py | 10 +++ tests/integrations/boto3/aws_mock.py | 33 ++++++++ tests/integrations/boto3/s3_list.xml | 2 + tests/integrations/boto3/test_s3.py | 85 +++++++++++++++++++ tox.ini | 7 ++ 7 files changed, 259 insertions(+) create mode 100644 sentry_sdk/integrations/boto3.py create mode 100644 tests/integrations/boto3/__init__.py create mode 100644 tests/integrations/boto3/aws_mock.py create mode 100644 tests/integrations/boto3/s3_list.xml create mode 100644 tests/integrations/boto3/test_s3.py diff --git a/sentry_sdk/integrations/__init__.py b/sentry_sdk/integrations/__init__.py index 3f0548ab63..777c363e14 100644 --- a/sentry_sdk/integrations/__init__.py +++ b/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", ) diff --git a/sentry_sdk/integrations/boto3.py b/sentry_sdk/integrations/boto3.py new file mode 100644 index 0000000000..573a6248bd --- /dev/null +++ b/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) diff --git a/tests/integrations/boto3/__init__.py b/tests/integrations/boto3/__init__.py new file mode 100644 index 0000000000..09738c40c7 --- /dev/null +++ b/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() diff --git a/tests/integrations/boto3/aws_mock.py b/tests/integrations/boto3/aws_mock.py new file mode 100644 index 0000000000..84ff23f466 --- /dev/null +++ b/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), + ) diff --git a/tests/integrations/boto3/s3_list.xml b/tests/integrations/boto3/s3_list.xml new file mode 100644 index 0000000000..10d5b16340 --- /dev/null +++ b/tests/integrations/boto3/s3_list.xml @@ -0,0 +1,2 @@ + +marshalls-furious-bucket1000urlfalsefoo.txt2020-10-24T00:13:39.000Z"a895ba674b4abd01b5d67cfd7074b827"2064537bef397f7e536914d1ff1bbdb105ed90bcfd06269456bf4a06c6e2e54564daf7STANDARDbar.txt2020-10-02T15:15:20.000Z"a895ba674b4abd01b5d67cfd7074b827"2064537bef397f7e536914d1ff1bbdb105ed90bcfd06269456bf4a06c6e2e54564daf7STANDARD diff --git a/tests/integrations/boto3/test_s3.py b/tests/integrations/boto3/test_s3.py new file mode 100644 index 0000000000..67376b55d4 --- /dev/null +++ b/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" diff --git a/tox.ini b/tox.ini index 98bfaf9a4d..4260c546cc 100644 --- a/tox.ini +++ b/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 =