From 679461137effed2a712c117fcd58bea30282e1dd Mon Sep 17 00:00:00 2001 From: Charlie Luo Date: Sun, 7 Sep 2025 16:09:57 -0700 Subject: [PATCH 1/4] add quota allowance to clickhouse errors --- snuba/web/db_query.py | 3 +++ tests/test_snql_api.py | 41 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py index 732b3fb9633..ae562b006ee 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -480,6 +480,9 @@ def _raw_query( calculated_cause = RateLimitExceeded( "Query scanned more than the allocated amount of bytes" ) + # Since we overwrite the original cause with a new error, we need to + # manually add the quota_allowance attribute to the new exception + calculated_cause.quota_allowance = stats["quota_allowance"] with configure_scope() as scope: fingerprint = ["{{default}}", str(cause.code), dataset_name] diff --git a/tests/test_snql_api.py b/tests/test_snql_api.py index 46f6cf904e8..ef274239a42 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", From 756ff2495830c790982e287d9cec3101a947a731 Mon Sep 17 00:00:00 2001 From: Charlie Luo Date: Sun, 14 Sep 2025 12:22:33 -0700 Subject: [PATCH 2/4] fix type --- snuba/web/db_query.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py index ae562b006ee..72b4a876783 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -477,12 +477,9 @@ def _raw_query( error_code = cause.code 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" - ) # Since we overwrite the original cause with a new error, we need to # manually add the quota_allowance attribute to the new exception - calculated_cause.quota_allowance = stats["quota_allowance"] + cause.extra_data["quota_allowance"] = stats["quota_allowance"] with configure_scope() as scope: fingerprint = ["{{default}}", str(cause.code), dataset_name] From 64cc915b575adc74723632414ef2ea51de7a8052 Mon Sep 17 00:00:00 2001 From: Charlie Luo Date: Wed, 14 Jan 2026 12:30:38 -0800 Subject: [PATCH 3/4] use cast --- snuba/state/rate_limit.py | 6 +++++- snuba/web/db_query.py | 7 ++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/snuba/state/rate_limit.py b/snuba/state/rate_limit.py index 6c2f64c8bf4..614930d8794 100644 --- a/snuba/state/rate_limit.py +++ b/snuba/state/rate_limit.py @@ -8,7 +8,7 @@ from contextlib import AbstractContextManager, ExitStack, contextmanager from dataclasses import dataclass from types import TracebackType -from typing import Any +from typing import Any, cast from typing import ChainMap as TypingChainMap from typing import Iterator, MutableMapping, Optional, Sequence, Type @@ -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 72b4a876783..e27a31183f5 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -477,9 +477,10 @@ def _raw_query( error_code = cause.code status = get_query_status_from_error_codes(error_code) if error_code == ErrorCodes.TOO_MANY_BYTES: - # Since we overwrite the original cause with a new error, we need to - # manually add the quota_allowance attribute to the new exception - cause.extra_data["quota_allowance"] = stats["quota_allowance"] + calculated_cause = RateLimitExceeded( + "Query scanned more than the allocated amount of bytes", + quota_allowance=stats["quota_allowance"], + ) with configure_scope() as scope: fingerprint = ["{{default}}", str(cause.code), dataset_name] From 76435b79ae3790897ae7a0a422ada34e19e409f1 Mon Sep 17 00:00:00 2001 From: "getsantry[bot]" <66042841+getsantry[bot]@users.noreply.github.com> Date: Wed, 14 Jan 2026 20:31:47 +0000 Subject: [PATCH 4/4] [getsentry/action-github-commit] Auto commit --- snuba/state/rate_limit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/snuba/state/rate_limit.py b/snuba/state/rate_limit.py index 614930d8794..7d4cdd1b960 100644 --- a/snuba/state/rate_limit.py +++ b/snuba/state/rate_limit.py @@ -8,9 +8,9 @@ from contextlib import AbstractContextManager, ExitStack, contextmanager from dataclasses import dataclass from types import TracebackType -from typing import Any, cast +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