From 2e0f2cf82b6a39fe4376e3aad7cb8a8a6e7ff9d2 Mon Sep 17 00:00:00 2001
From: Rachel Chen <rachelchen@PL6VFX9HP4.attlocal.net>
Date: Mon, 24 Jun 2024 11:13:51 -0700
Subject: [PATCH] single dictionary

---
 snuba/query/allocation_policies/__init__.py | 41 ++++++++-------------
 snuba/web/db_query.py                       |  2 +-
 tests/test_snql_api.py                      | 23 ++++++------
 3 files changed, 29 insertions(+), 37 deletions(-)

diff --git a/snuba/query/allocation_policies/__init__.py b/snuba/query/allocation_policies/__init__.py
index fc4acf4149c..96470e89156 100644
--- a/snuba/query/allocation_policies/__init__.py
+++ b/snuba/query/allocation_policies/__init__.py
@@ -158,7 +158,7 @@ class AllocationPolicyViolations(SerializableException):
     """
 
     def __str__(self) -> str:
-        return f"{self.message}, details: {self.violations | self.summary}"
+        return f"{self.message}, details: {self.violations}, summary: {self.summary}"
 
     @property
     def violations(self) -> dict[str, dict[str, Any]]:
@@ -166,39 +166,30 @@ def violations(self) -> dict[str, dict[str, Any]]:
 
     @property
     def quota_allowance(self) -> dict[str, dict[str, Any]]:
-        return {
-            k: v
-            for k, v in self._quota_allowances_and_summary.items()
-            if k != "summary"
-        }
+        return cast(
+            dict[str, dict[str, Any]], self.extra_data.get("quota_allowances", {})
+        )
 
     @property
     def summary(self) -> dict[str, Any]:
-        return {
-            k: v
-            for k, v in self._quota_allowances_and_summary.items()
-            if k == "summary"
-        }
-
-    @property
-    def _quota_allowances_and_summary(self) -> dict[str, dict[str, Any]]:
-        return cast(
-            dict[str, dict[str, Any]],
-            self.extra_data.get("quota_allowances_and_summary", {}),
-        )
+        return cast(dict[str, Any], self.extra_data.get("summary"))
 
     @classmethod
     def from_args(
-        cls, quota_allowances_and_summary: dict[str, Any]
+        cls,
+        quota_allowances: dict[str, QuotaAllowance],
+        summary: dict[str, Any],
     ) -> "AllocationPolicyViolations":
-        obj = cls(
+        quota_allowances_dict = {
+            key: quota_allowance.to_dict()
+            for key, quota_allowance in quota_allowances.items()
+        }
+
+        return cls(
             "Query on could not be run due to allocation policies",
-            quota_allowances_and_summary={
-                key: quota_allowance
-                for key, quota_allowance in quota_allowances_and_summary.items()
-            },
+            quota_allowances=quota_allowances_dict,
+            summary=summary,
         )
-        return obj
 
 
 class AllocationPolicy(ABC, metaclass=RegisteredClass):
diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py
index c343b9d79e7..67d52d86946 100644
--- a/snuba/web/db_query.py
+++ b/snuba/web/db_query.py
@@ -900,7 +900,7 @@ def _apply_allocation_policies_quota(
         stats["quota_allowance"]["summary"] = summary
 
         if not can_run:
-            raise AllocationPolicyViolations.from_args(stats["quota_allowance"])
+            raise AllocationPolicyViolations.from_args(quota_allowances, summary)
         # Before allocation policies were a thing, the query pipeline would apply
         # thread limits in a query processor. That is not necessary if there
         # is an allocation_policy in place but nobody has removed that code yet.
diff --git a/tests/test_snql_api.py b/tests/test_snql_api.py
index bd016c3689d..a3cf572dbb1 100644
--- a/tests/test_snql_api.py
+++ b/tests/test_snql_api.py
@@ -1311,21 +1311,22 @@ def test_allocation_policy_violation(self) -> None:
                     "quota_unit": NO_UNITS,
                     "suggestion": NO_SUGGESTION,
                 },
-                "summary": {
-                    "threads_used": 0,
-                    "rejected_by": {
-                        "policy": "RejectAllocationPolicy123",
-                        "quota_used": 0,
-                        "quota_unit": "no_units",
-                        "suggestion": "no_suggestion",
-                        "rejection_threshold": 1000000000000,
-                    },
-                    "throttled_by": {},
+            }
+            summary = {
+                "threads_used": 0,
+                "rejected_by": {
+                    "policy": "RejectAllocationPolicy123",
+                    "quota_used": 0,
+                    "quota_unit": "no_units",
+                    "suggestion": "no_suggestion",
+                    "rejection_threshold": 1000000000000,
                 },
+                "throttled_by": {},
             }
+
             assert (
                 response.json["error"]["message"]
-                == f"Query on could not be run due to allocation policies, details: {details}"
+                == f"Query on could not be run due to allocation policies, details: {details}, summary: {summary}"
             )
 
     def test_tags_key_column(self) -> None: