diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py index d02b3e1110..22fbfb65d0 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -847,6 +847,26 @@ def _add_quota_info( quota_info["throttle_threshold"] = quota_allowance.throttle_threshold +def _populate_query_status( + summary: dict[str, Any], + rejection_quota_and_policy: Optional[_QuotaAndPolicy], + throttle_quota_and_policy: Optional[_QuotaAndPolicy], +) -> None: + is_successful = "is_successful" + is_rejected = "is_rejected" + is_throttled = "is_throttled" + summary[is_successful] = True + summary[is_rejected] = False + summary[is_throttled] = False + + if rejection_quota_and_policy: + summary[is_successful] = False + summary[is_rejected] = True + if throttle_quota_and_policy: + summary[is_successful] = False + summary[is_throttled] = True + + def _apply_allocation_policies_quota( query_settings: QuerySettings, attribution_info: AttributionInfo, @@ -908,6 +928,9 @@ def _apply_allocation_policies_quota( summary: dict[str, Any] = {} summary["threads_used"] = min_threads_across_policies + _populate_query_status( + summary, rejection_quota_and_policy, throttle_quota_and_policy + ) _add_quota_info(summary, _REJECTED_BY, rejection_quota_and_policy) _add_quota_info(summary, _THROTTLED_BY, throttle_quota_and_policy) stats["quota_allowance"]["summary"] = summary diff --git a/tests/test_snql_api.py b/tests/test_snql_api.py index 4d518e2c94..9a49144b29 100644 --- a/tests/test_snql_api.py +++ b/tests/test_snql_api.py @@ -1348,6 +1348,9 @@ def test_allocation_policy_violation(self) -> None: }, "summary": { "threads_used": 0, + "is_successful": False, + "is_rejected": True, + "is_throttled": False, "rejected_by": { "policy": "RejectAllocationPolicy123", "quota_used": 0, diff --git a/tests/web/test_db_query.py b/tests/web/test_db_query.py index 63d4781678..92beb10373 100644 --- a/tests/web/test_db_query.py +++ b/tests/web/test_db_query.py @@ -306,6 +306,9 @@ def test_db_query_success() -> None: assert stats["quota_allowance"] == { "summary": { "threads_used": 5, + "is_successful": False, + "is_rejected": False, + "is_throttled": True, "rejected_by": {}, "throttled_by": { "policy": "BytesScannedRejectingPolicy", @@ -584,6 +587,9 @@ def __init__(self, max_threads: int, policy_name: str) -> None: }, "summary": { "threads_used": 1, + "is_successful": False, + "is_rejected": False, + "is_throttled": True, "rejected_by": {}, "throttled_by": { "policy": "ThrottleAllocationPolicy1", @@ -657,6 +663,9 @@ def _update_quota_balance( assert stats["quota_allowance"] == { "summary": { "threads_used": 0, + "is_successful": False, + "is_rejected": True, + "is_throttled": True, "rejected_by": { "policy": "RejectAllocationPolicy", "rejection_threshold": MAX_THRESHOLD, @@ -897,6 +906,9 @@ def _run_query() -> None: assert e.value.extra["stats"]["quota_allowance"] == { "summary": { "threads_used": 0, + "is_successful": False, + "is_rejected": True, + "is_throttled": False, "rejected_by": { "policy": "CountQueryPolicy", "rejection_threshold": MAX_QUERIES_TO_RUN,