diff --git a/sentry_sdk/transport.py b/sentry_sdk/transport.py index 60ab611c54..6d6a1c1f91 100644 --- a/sentry_sdk/transport.py +++ b/sentry_sdk/transport.py @@ -15,17 +15,22 @@ from sentry_sdk._types import MYPY if MYPY: - from typing import Type from typing import Any - from typing import Optional + from typing import Callable from typing import Dict + from typing import Iterable + from typing import Optional + from typing import Tuple + from typing import Type from typing import Union - from typing import Callable + from urllib3.poolmanager import PoolManager # type: ignore from urllib3.poolmanager import ProxyManager from sentry_sdk._types import Event + DataCategory = Optional[str] + try: from urllib.request import getproxies except ImportError: @@ -94,6 +99,21 @@ def __del__(self): pass +def _parse_rate_limits(header, now=None): + # type: (Any, Optional[datetime]) -> Iterable[Tuple[DataCategory, datetime]] + if now is None: + now = datetime.utcnow() + + for limit in header.split(","): + try: + retry_after, categories, _ = limit.strip().split(":", 2) + retry_after = now + timedelta(seconds=int(retry_after)) + for category in categories and categories.split(";") or (None,): + yield category, retry_after + except (LookupError, ValueError): + continue + + class HttpTransport(Transport): """The default HTTP transport.""" @@ -107,7 +127,7 @@ def __init__( assert self.parsed_dsn is not None self._worker = BackgroundWorker() self._auth = self.parsed_dsn.to_auth("sentry.python/%s" % VERSION) - self._disabled_until = {} # type: Dict[Any, datetime] + self._disabled_until = {} # type: Dict[DataCategory, datetime] self._retry = urllib3.util.Retry() self.options = options @@ -129,16 +149,7 @@ def _update_rate_limits(self, response): # no matter of the status code to update our internal rate limits. header = response.headers.get("x-sentry-rate-limit") if header: - for limit in header.split(","): - try: - retry_after, categories, _ = limit.strip().split(":", 2) - retry_after = datetime.utcnow() + timedelta( - seconds=int(retry_after) - ) - for category in categories.split(";") or (None,): - self._disabled_until[category] = retry_after - except (LookupError, ValueError): - continue + self._disabled_until.update(_parse_rate_limits(header)) # old sentries only communicate global rate limit hits via the # retry-after header on 429. This header can also be emitted on new diff --git a/tests/test_transport.py b/tests/test_transport.py index 00cdc6c42e..398ff0a6da 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -2,11 +2,12 @@ import logging import pickle -from datetime import datetime +from datetime import datetime, timedelta import pytest from sentry_sdk import Hub, Client, add_breadcrumb, capture_message +from sentry_sdk.transport import _parse_rate_limits @pytest.fixture(params=[True, False]) @@ -54,3 +55,115 @@ def test_transport_works( assert httpserver.requests assert any("Sending event" in record.msg for record in caplog.records) == debug + + +NOW = datetime(2014, 6, 2) + + +@pytest.mark.parametrize( + "input,expected", + [ + # Invalid rate limits + ("", {}), + ("invalid", {}), + (",,,", {}), + ( + "42::organization, invalid, 4711:foobar;transaction;security:project", + { + None: NOW + timedelta(seconds=42), + "transaction": NOW + timedelta(seconds=4711), + "security": NOW + timedelta(seconds=4711), + # Unknown data categories + "foobar": NOW + timedelta(seconds=4711), + }, + ), + ( + "4711:foobar;;transaction:organization", + { + "transaction": NOW + timedelta(seconds=4711), + # Unknown data categories + "foobar": NOW + timedelta(seconds=4711), + "": NOW + timedelta(seconds=4711), + }, + ), + ], +) +def test_parse_rate_limits(input, expected): + assert dict(_parse_rate_limits(input, now=NOW)) == expected + + +def test_simple_rate_limits(httpserver, capsys, caplog): + client = Client(dsn="http://foobar@{}/123".format(httpserver.url[len("http://") :])) + httpserver.serve_content("no", 429, headers={"Retry-After": "4"}) + + client.capture_event({"type": "transaction"}) + client.flush() + + assert len(httpserver.requests) == 1 + del httpserver.requests[:] + + assert set(client.transport._disabled_until) == set([None]) + + client.capture_event({"type": "transaction"}) + client.capture_event({"type": "event"}) + client.flush() + + assert not httpserver.requests + + +@pytest.mark.parametrize("response_code", [200, 429]) +def test_data_category_limits(httpserver, capsys, caplog, response_code): + client = Client( + dict(dsn="http://foobar@{}/123".format(httpserver.url[len("http://") :])) + ) + httpserver.serve_content( + "hm", + response_code, + headers={"X-Sentry-Rate-Limit": "4711:transaction:organization"}, + ) + + client.capture_event({"type": "transaction"}) + client.flush() + + assert len(httpserver.requests) == 1 + del httpserver.requests[:] + + assert set(client.transport._disabled_until) == set(["transaction"]) + + client.transport.capture_event({"type": "transaction"}) + client.transport.capture_event({"type": "transaction"}) + client.flush() + + assert not httpserver.requests + + client.capture_event({"type": "event"}) + client.flush() + + assert len(httpserver.requests) == 1 + + +@pytest.mark.parametrize("response_code", [200, 429]) +def test_complex_limits_without_data_category( + httpserver, capsys, caplog, response_code +): + client = Client( + dict(dsn="http://foobar@{}/123".format(httpserver.url[len("http://") :])) + ) + httpserver.serve_content( + "hm", response_code, headers={"X-Sentry-Rate-Limit": "4711::organization"}, + ) + + client.capture_event({"type": "transaction"}) + client.flush() + + assert len(httpserver.requests) == 1 + del httpserver.requests[:] + + assert set(client.transport._disabled_until) == set([None]) + + client.transport.capture_event({"type": "transaction"}) + client.transport.capture_event({"type": "transaction"}) + client.capture_event({"type": "event"}) + client.flush() + + assert len(httpserver.requests) == 0