diff --git a/snuba/state/rate_limit.py b/snuba/state/rate_limit.py index 6c2f64c8bf..7d4cdd1b96 100644 --- a/snuba/state/rate_limit.py +++ b/snuba/state/rate_limit.py @@ -10,7 +10,7 @@ from types import TracebackType from typing import Any from typing import ChainMap as TypingChainMap -from typing import Iterator, MutableMapping, Optional, Sequence, Type +from typing import Iterator, MutableMapping, Optional, Sequence, Type, cast from snuba import environment, state from snuba.redis import RedisClientKey, get_redis_client @@ -75,6 +75,10 @@ class RateLimitExceeded(SerializableException): additional parameters which are provided when the exception is raised. """ + @property + def quota_allowance(self) -> dict[str, Any]: + return cast(dict[str, Any], self.extra_data.get("quota_allowance", {})) + @dataclass(frozen=True) class RateLimitStats: diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py index 732b3fb963..e27a31183f 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -478,7 +478,8 @@ def _raw_query( status = get_query_status_from_error_codes(error_code) if error_code == ErrorCodes.TOO_MANY_BYTES: calculated_cause = RateLimitExceeded( - "Query scanned more than the allocated amount of bytes" + "Query scanned more than the allocated amount of bytes", + quota_allowance=stats["quota_allowance"], ) with configure_scope() as scope: diff --git a/tests/test_snql_api.py b/tests/test_snql_api.py index 46f6cf904e..ef274239a4 100644 --- a/tests/test_snql_api.py +++ b/tests/test_snql_api.py @@ -1360,6 +1360,45 @@ def test_allocation_policy_max_bytes_to_read(self) -> None: == "Query scanned more than the allocated amount of bytes" ) + expected_quota_allowance = { + "details": { + "MaxBytesPolicy123": { + "can_run": True, + "max_threads": 0, + "max_bytes_to_read": 1, + "explanation": { + "storage_key": "doesntmatter", + }, + "is_throttled": True, + "throttle_threshold": MAX_THRESHOLD, + "rejection_threshold": MAX_THRESHOLD, + "quota_used": 0, + "quota_unit": NO_UNITS, + "suggestion": NO_SUGGESTION, + } + }, + "summary": { + "threads_used": 0, + "max_bytes_to_read": 1, + "is_successful": False, + "is_rejected": False, + "is_throttled": True, + "rejection_storage_key": None, + "throttle_storage_key": "doesntmatter", + "rejected_by": {}, + "throttled_by": { + "policy": "MaxBytesPolicy123", + "quota_used": 0, + "quota_unit": NO_UNITS, + "suggestion": NO_SUGGESTION, + "storage_key": "doesntmatter", + "throttle_threshold": MAX_THRESHOLD, + }, + }, + } + + assert response.json["quota_allowance"] == expected_quota_allowance + def test_allocation_policy_violation(self) -> None: with patch( "snuba.web.db_query._get_allocation_policies", @@ -1427,6 +1466,8 @@ def test_allocation_policy_violation(self) -> None: == f"Query on could not be run due to allocation policies, info: {info}" ) + assert response.json["quota_allowance"] == info + def test_tags_key_column(self) -> None: response = self.post( "/events/snql",