From 0670326e5d3a2c237c66c1b618cb9d5be6ed1141 Mon Sep 17 00:00:00 2001 From: Charlie Luo Date: Sun, 7 Sep 2025 16:09:57 -0700 Subject: [PATCH 1/2] 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 1cb062f839a..62226de195c 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -486,6 +486,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 87bc454f27b..a4d347f4ae0 100644 --- a/tests/test_snql_api.py +++ b/tests/test_snql_api.py @@ -1370,6 +1370,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", @@ -1437,6 +1476,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 eb5e328f77549aaf1e47a0e65cada80022fcb43b Mon Sep 17 00:00:00 2001 From: Charlie Luo Date: Sun, 14 Sep 2025 12:22:33 -0700 Subject: [PATCH 2/2] 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 62226de195c..cfb3917d025 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -483,12 +483,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]