feat(allocation_policies): add org allowlist to bypass all rate limits

snuba/query/allocation_policies/__init__.py

@@ -18,6 +18,7 @@
     logger,
 )
 from snuba.datasets.storages.storage_key import StorageKey
+from snuba.state import get_config as get_runtime_config
 from snuba.utils.metrics.wrapper import MetricsWrapper
 from snuba.utils.registered_class import import_submodules_in_directory
 from snuba.utils.serializable_exception import JsonSerializable, SerializableException
@@ -26,6 +27,13 @@
 CAPMAN_PREFIX = "capman"
 CAPMAN_HASH = "capman"
 
+# Global runtime config (visible in the runtime config admin UI, NOT Capacity
+# Management) holding a comma-separated list of organization ids that are exempt
+# from ALL allocation policies (rate limits), e.g. "1,42,1337". It is global
+# rather than per-policy because it is meant to exempt an org from every policy
+# at once. See AllocationPolicy.is_rate_limit_bypassed.
+ORG_RATE_LIMIT_BYPASS_CONFIG = "organization_rate_limit_bypass_allowlist"
+
 IS_ACTIVE = "is_active"
 IS_ENFORCED = "is_enforced"
 MAX_THREADS = "max_threads"
@@ -424,6 +432,30 @@ def __eq__(self, other: Any) -> bool:
     def is_cross_org_query(self, tenant_ids: dict[str, str | int]) -> bool:
         return bool(tenant_ids.get("cross_org_query", False))
 
+    def is_rate_limit_bypassed(self, tenant_ids: dict[str, str | int]) -> bool:
+        """Returns True if the query's organization is in the global rate limit
+        bypass allowlist. Orgs in this allowlist skip ALL allocation policies
+        entirely, in the same way an inactive policy is skipped.
+
+        The allowlist is a single global runtime config shared across every
+        policy (see ORG_RATE_LIMIT_BYPASS_CONFIG), since it is meant to exempt an
+        org from all rate limits at once. It is stored as a comma-separated list
+        of organization ids, e.g. "1,42,1337".
+        """
+        if not tenant_ids:
+            return False
+        org_id = tenant_ids.get("organization_id")
+        if org_id is None:
+            return False
+        # Use `is None` rather than a falsy check: runtime config coerces types,
+        # so a lone org id of `0` would be stored as the integer 0 and a `not
+        # allowlist` check would wrongly treat it as unset.
+        allowlist = get_runtime_config(ORG_RATE_LIMIT_BYPASS_CONFIG)
+        if allowlist is None:
+            return False
+        bypassed_orgs = {part.strip() for part in str(allowlist).split(",") if part.strip()}
+        return str(org_id) in bypassed_orgs
+
     @classmethod
     def from_kwargs(cls, **kwargs: str) -> AllocationPolicy:
         required_tenant_types = kwargs.pop("required_tenant_types", None)
@@ -457,7 +489,7 @@ def get_quota_allowance(
             for t, tid in tenant_ids.items():
                 span.set_data(f"tenant_ids.{t}", str(tid))
             try:
-                if not self.is_active:
+                if not self.is_active or self.is_rate_limit_bypassed(tenant_ids):
                     allowance = QuotaAllowance(
                         can_run=True,
                         max_threads=self.max_threads,
@@ -547,7 +579,7 @@ def update_quota_balance(
         result_or_error: QueryResultOrError,
     ) -> None:
         try:
-            if not self.is_active:
+            if not self.is_active or self.is_rate_limit_bypassed(tenant_ids):
                 return None
             return self._update_quota_balance(tenant_ids, query_id, result_or_error)
         except InvalidTenantsForAllocationPolicy:

tests/query/allocation_policies/test_allocation_policy_base.py

@@ -12,6 +12,7 @@
     MAX_THRESHOLD,
     NO_SUGGESTION,
     NO_UNITS,
+    ORG_RATE_LIMIT_BYPASS_CONFIG,
     AllocationPolicy,
     AllocationPolicyConfig,
     InvalidTenantsForAllocationPolicy,
@@ -524,6 +525,94 @@ def test_is_not_enforced() -> None:
     assert throttled_metrics[1].tags["is_enforced"] == "False"
 
 
+@pytest.mark.redis_db
+def test_org_rate_limit_bypass_allowlist() -> None:
+    MAX_THREADS = 100
+    reject_policy = RejectingEverythingAllocationPolicy(
+        StorageKey("some_storage"),
+        [],
+        {"is_active": 1, "is_enforced": 1, "max_threads": MAX_THREADS},
+    )
+    allowlisted_tenant: dict[str, int | str] = {
+        "organization_id": 123,
+        "referrer": "some_referrer",
+    }
+    other_tenant: dict[str, int | str] = {
+        "organization_id": 999,
+        "referrer": "some_referrer",
+    }
+
+    # By default the org is rejected by the (enforced) policy.
+    assert not reject_policy.get_quota_allowance(allowlisted_tenant, "deadbeef").can_run
+
+    # Add the org to the global bypass allowlist.
+    set_config(ORG_RATE_LIMIT_BYPASS_CONFIG, "123,456")
+
+    # Allowlisted org bypasses the policy entirely (passes through with full threads).
+    allowance = reject_policy.get_quota_allowance(allowlisted_tenant, "deadbeef")
+    assert allowance.can_run
+    assert allowance.max_threads == MAX_THREADS
+
+    # An org not in the allowlist is still rejected.
+    assert not reject_policy.get_quota_allowance(other_tenant, "deadbeef").can_run
+
+
+@pytest.mark.redis_db
+def test_org_rate_limit_bypass_skips_update_quota_balance() -> None:
+    # BadlyWrittenAllocationPolicy raises in both _get_quota_allowance and
+    # _update_quota_balance, so a successful call proves those private methods
+    # were never invoked (i.e. the org bypassed the policy).
+    policy = BadlyWrittenAllocationPolicy(
+        StorageKey("some_storage"),
+        [],
+        {"is_active": 1, "is_enforced": 1},
+    )
+    tenant_ids: dict[str, int | str] = {
+        "organization_id": 123,
+        "referrer": "some_referrer",
+    }
+    result_or_error = QueryResultOrError(
+        query_result=QueryResult(
+            result={"profile": {"bytes": 420}},
+            extra={"stats": {}, "sql": "", "experiments": {}},
+        ),
+        error=None,
+    )
+
+    set_config(ORG_RATE_LIMIT_BYPASS_CONFIG, "123")
+
+    # Neither the buggy _get_quota_allowance nor _update_quota_balance is reached.
+    assert policy.get_quota_allowance(tenant_ids, "deadbeef").can_run
+    policy.update_quota_balance(tenant_ids, "deadbeef", result_or_error)
+
+
+@pytest.mark.redis_db
+def test_org_rate_limit_bypass_no_org_id() -> None:
+    reject_policy = RejectingEverythingAllocationPolicy(
+        StorageKey("some_storage"),
+        [],
+        {"is_active": 1, "is_enforced": 1},
+    )
+    set_config(ORG_RATE_LIMIT_BYPASS_CONFIG, "123")
+    # tenant_ids without an organization_id should never be bypassed.
+    assert not reject_policy.get_quota_allowance({"referrer": "some_referrer"}, "deadbeef").can_run
+
+
+@pytest.mark.redis_db
+def test_org_rate_limit_bypass_zero_org_id() -> None:
+    # A lone org id of 0 is coerced to the integer 0 by runtime config; make
+    # sure it is still honored (i.e. not treated as an unset allowlist).
+    reject_policy = RejectingEverythingAllocationPolicy(
+        StorageKey("some_storage"),
+        [],
+        {"is_active": 1, "is_enforced": 1},
+    )
+    set_config(ORG_RATE_LIMIT_BYPASS_CONFIG, "0")
+    assert reject_policy.get_quota_allowance(
+        {"organization_id": 0, "referrer": "some_referrer"}, "deadbeef"
+    ).can_run
+
+
 @pytest.mark.redis_db
 def test_configs_with_delimiter_values() -> None:
     # test that configs with dots can be stored and read