test

getsentry · Jun 17, 2024 · de4acc5 · de4acc5
1 parent a79dc39
commit de4acc5
Show file tree

Hide file tree

Showing 6 changed files with 144 additions and 6 deletions.
diff --git a/snuba/query/allocation_policies/__init__.py b/snuba/query/allocation_policies/__init__.py
@@ -764,7 +764,7 @@ def get_quota_allowance(
                 },
             )
         if not self.is_enforced:
-            return QuotaAllowance(True, self.max_threads, {})
+            return QuotaAllowance(True, self.max_threads, {}, True)
         # make sure we always know which storage key we rejected a query from
         allowance.explanation["storage_key"] = str(self._storage_key)
         return allowance

diff --git a/tests/admin/test_api.py b/tests/admin/test_api.py
@@ -427,7 +427,7 @@ def _additional_config_definitions(self) -> list[AllocationPolicyConfig]:
         def _get_quota_allowance(
             self, tenant_ids: dict[str, str | int], query_id: str
         ) -> QuotaAllowance:
-            return QuotaAllowance(True, 1, {})
+            return QuotaAllowance(True, 1, {}, False)
 
         def _update_quota_balance(
             self,
@@ -437,6 +437,21 @@ def _update_quota_balance(
         ) -> None:
             pass
 
+        def get_throttle_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_rejection_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_used(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_units(self) -> str:
+            return "No units"
+
+        def get_suggestion(self) -> str:
+            return "No suggestion"
+
     def mock_get_policies() -> list[AllocationPolicy]:
         policy = FakePolicy(StorageKey("nothing"), [], {})
         policy.set_config_value("fake_optional_config", 10, {"org_id": 10})

diff --git a/tests/pipeline/test_execution_stage.py b/tests/pipeline/test_execution_stage.py
@@ -41,7 +41,9 @@ def _additional_config_definitions(self) -> list[AllocationPolicyConfig]:
     def _get_quota_allowance(
         self, tenant_ids: dict[str, str | int], query_id: str
     ) -> QuotaAllowance:
-        return QuotaAllowance(can_run=True, max_threads=1, explanation={})
+        return QuotaAllowance(
+            can_run=True, max_threads=1, explanation={}, is_throttled=False
+        )
 
     def _update_quota_balance(
         self,
@@ -51,6 +53,21 @@ def _update_quota_balance(
     ) -> None:
         self.did_update = True
 
+    def get_throttle_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+        return -1
+
+    def get_rejection_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+        return -1
+
+    def get_quota_used(self, tenant_ids: dict[str, str | int]) -> int:
+        return -1
+
+    def get_quota_units(self) -> str:
+        return "No units"
+
+    def get_suggestion(self) -> str:
+        return "No suggestion"
+
 
 def get_fake_metadata() -> SnubaQueryMetadata:
     return SnubaQueryMetadata(

diff --git a/tests/query/allocation_policies/test_allocation_policy_base.py b/tests/query/allocation_policies/test_allocation_policy_base.py
@@ -60,14 +60,18 @@ class RejectingEverythingAllocationPolicy(PassthroughPolicy):
     def _get_quota_allowance(
         self, tenant_ids: dict[str, str | int], query_id: str
     ) -> QuotaAllowance:
-        return QuotaAllowance(can_run=False, max_threads=10, explanation={})
+        return QuotaAllowance(
+            can_run=False, max_threads=10, explanation={}, is_throttled=False
+        )
 
 
 class ThrottleEverythingAllocationPolicy(PassthroughPolicy):
     def _get_quota_allowance(
         self, tenant_ids: dict[str, str | int], query_id: str
     ) -> QuotaAllowance:
-        return QuotaAllowance(can_run=True, max_threads=1, explanation={})
+        return QuotaAllowance(
+            can_run=True, max_threads=1, explanation={}, is_throttled=True
+        )
 
 
 class BadlyWrittenAllocationPolicy(PassthroughPolicy):
@@ -192,6 +196,21 @@ def _update_quota_balance(
     ) -> None:
         pass
 
+    def get_throttle_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+        return -1
+
+    def get_rejection_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+        return -1
+
+    def get_quota_used(self, tenant_ids: dict[str, str | int]) -> int:
+        return -1
+
+    def get_quota_units(self) -> str:
+        return "No units"
+
+    def get_suggestion(self) -> str:
+        return "No suggestion"
+
 
 class TestAllocationPolicyLogs(TestCase):
     @pytest.mark.redis_db

diff --git a/tests/test_snql_api.py b/tests/test_snql_api.py
@@ -40,7 +40,7 @@ def _get_quota_allowance(
             can_run=False,
             max_threads=0,
             explanation={"reason": "policy rejects all queries"},
-            is_throttled=True,
+            is_throttled=False,
         )
 
     def _update_quota_balance(
@@ -51,6 +51,21 @@ def _update_quota_balance(
     ) -> None:
         return
 
+    def get_throttle_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+        return -1
+
+    def get_rejection_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+        return -1
+
+    def get_quota_used(self, tenant_ids: dict[str, str | int]) -> int:
+        return -1
+
+    def get_quota_units(self) -> str:
+        return "No units"
+
+    def get_suggestion(self) -> str:
+        return "No suggestion"
+
 
 @pytest.mark.clickhouse_db
 @pytest.mark.redis_db

diff --git a/tests/web/test_db_query.py b/tests/web/test_db_query.py
@@ -272,6 +272,7 @@ def test_db_query_success() -> None:
                 "referrer": "something",
                 "storage_key": "StorageKey.ERRORS_RO",
             },
+            "is_throttled": False,
         },
         "ConcurrentRateLimitAllocationPolicy": {
             "can_run": True,
@@ -281,6 +282,7 @@ def test_db_query_success() -> None:
                 "overrides": {},
                 "storage_key": "StorageKey.ERRORS_RO",
             },
+            "is_throttled": False,
         },
         "BytesScannedRejectingPolicy": {
             "can_run": True,
@@ -289,6 +291,7 @@ def test_db_query_success() -> None:
                 "reason": "within_limit but throttled",
                 "storage_key": "StorageKey.ERRORS_RO",
             },
+            "is_throttled": True,
         },
         "CrossOrgQueryAllocationPolicy": {
             "can_run": True,
@@ -297,11 +300,13 @@ def test_db_query_success() -> None:
                 "reason": "pass_through",
                 "storage_key": "StorageKey.ERRORS_RO",
             },
+            "is_throttled": False,
         },
         "BytesScannedWindowAllocationPolicy": {
             "can_run": True,
             "max_threads": 10,
             "explanation": {"storage_key": "StorageKey.ERRORS_RO"},
+            "is_throttled": True,
         },
     }
 
@@ -415,6 +420,7 @@ def _get_quota_allowance(
                 can_run=False,
                 max_threads=0,
                 explanation={"reason": "policy rejects all queries"},
+                is_throttled=False,
             )
 
         def _update_quota_balance(
@@ -427,6 +433,21 @@ def _update_quota_balance(
             update_called = True
             return
 
+        def get_throttle_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_rejection_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_used(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_units(self) -> str:
+            return "No units"
+
+        def get_suggestion(self) -> str:
+            return "No suggestion"
+
     with mock.patch(
         "snuba.web.db_query._get_allocation_policies",
         return_value=[
@@ -457,6 +478,7 @@ def _update_quota_balance(
                     "storage_key": "StorageKey.DOESNTMATTER",
                 },
                 "max_threads": 0,
+                "is_throttled": False,
             }
         }
         # extra data contains policy failure information
@@ -491,6 +513,7 @@ def _get_quota_allowance(
                 can_run=True,
                 max_threads=POLICY_THREADS,
                 explanation={"reason": "Throttle everything!"},
+                is_throttled=False,
             )
 
         def _update_quota_balance(
@@ -501,6 +524,21 @@ def _update_quota_balance(
         ) -> None:
             return
 
+        def get_throttle_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_rejection_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_used(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_units(self) -> str:
+            return "No units"
+
+        def get_suggestion(self) -> str:
+            return "No suggestion"
+
     class ThreadLimitPolicyDuplicate(ThreadLimitPolicy):
         def _get_quota_allowance(
             self, tenant_ids: dict[str, str | int], query_id: str
@@ -509,6 +547,7 @@ def _get_quota_allowance(
                 can_run=True,
                 max_threads=POLICY_THREADS + 1,
                 explanation={"reason": "Throttle everything!"},
+                is_throttled=False,
             )
 
     # Should limit to minimal threads across policies
@@ -563,6 +602,7 @@ def _get_quota_allowance(
                 can_run=can_run,
                 max_threads=0,
                 explanation={"reason": f"can only run {queries_run} queries!"},
+                is_throttled=False,
             )
 
         def _update_quota_balance(
@@ -574,6 +614,21 @@ def _update_quota_balance(
             nonlocal queries_run
             queries_run += 1
 
+        def get_throttle_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_rejection_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_used(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_units(self) -> str:
+            return "No units"
+
+        def get_suggestion(self) -> str:
+            return "No suggestion"
+
     queries_run_duplicate = 0
 
     class CountQueryPolicyDuplicate(AllocationPolicy):
@@ -592,6 +647,7 @@ def _get_quota_allowance(
                 explanation={
                     "reason": f"can only run {queries_run_duplicate} queries!"
                 },
+                is_throttled=False,
             )
 
         def _update_quota_balance(
@@ -603,6 +659,21 @@ def _update_quota_balance(
             nonlocal queries_run_duplicate
             queries_run_duplicate += 1
 
+        def get_throttle_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_rejection_threshold(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_used(self, tenant_ids: dict[str, str | int]) -> int:
+            return -1
+
+        def get_quota_units(self) -> str:
+            return "No units"
+
+        def get_suggestion(self) -> str:
+            return "No suggestion"
+
     # the first policy will error and short circuit the rest
     query, storage, attribution_info = _build_test_query(
         "count(distinct(project_id))",
@@ -643,6 +714,7 @@ def _run_query() -> None:
                 "reason": "can only run 2 queries!",
                 "storage_key": "StorageKey.DOESNTMATTER",
             },
+            "is_throttled": False,
         },
     }
     cause = e.value.__cause__