More comprehensive payload from CapMan (#6256)

https://getsentry.atlassian.net/browse/SNS-2852 Allocation policies are our mechanism for doing traffic management for Snuba queries. This PR adds more fields to the payload returned by Snuba, so we can filter more efficiently in Sentry traces Co-authored-by: Rachel Chen <[email protected]>
getsentry · Aug 29, 2024 · d703a97 · d703a97
1 parent 69b9152
commit d703a97
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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,