diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py index 99c37b512fc..9993b566f84 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -476,9 +476,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 + cause.extra_data["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",