From de4acc54c4269059f3a81144181f24b0bf2d795d Mon Sep 17 00:00:00 2001 From: Rachel Chen Date: Mon, 17 Jun 2024 13:05:14 -0700 Subject: [PATCH] test --- snuba/query/allocation_policies/__init__.py | 2 +- tests/admin/test_api.py | 17 ++++- tests/pipeline/test_execution_stage.py | 19 ++++- .../test_allocation_policy_base.py | 23 +++++- tests/test_snql_api.py | 17 ++++- tests/web/test_db_query.py | 72 +++++++++++++++++++ 6 files changed, 144 insertions(+), 6 deletions(-) diff --git a/snuba/query/allocation_policies/__init__.py b/snuba/query/allocation_policies/__init__.py index 3bdb2653ab..4e301e8f3b 100644 --- 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 index 58f38ea578..f5305e2c8b 100644 --- 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 index 2a290932b2..b908064e0c 100644 --- 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 index 53b46ddc2d..fcafc0d046 100644 --- 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 index a80e88b1aa..7df8ebf09e 100644 --- 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 index 53d861f5ea..6b9bac6fe6 100644 --- 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__