squash

Rachel Chen · Rachel Chen · commit 6731d49d7768 · 2025-09-23T14:56:36.000-07:00
mypy

uuid

max bytes to read

tests

move to util

query id

query id

combine recommendations

remove prchanges?

move get_max_bytes_to_read

AllocationPolicyViolations

allocation policy fixes and tests

remove prchagnes
diff --git a/snuba/datasets/configuration/events_analytics_platform/storages/eap_items.yaml b/snuba/datasets/configuration/events_analytics_platform/storages/eap_items.yaml
@@ -109,31 +109,6 @@ schema:
   local_table_name: eap_items_1_local
   dist_table_name: eap_items_1_dist
   partition_format: [date]
-allocation_policies:
-  - name: ConcurrentRateLimitAllocationPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - referrer
-        - project_id
-      default_config_overrides:
-        is_enforced: 0
-  - name: ReferrerGuardRailPolicy
-    args:
-      required_tenant_types:
-        - referrer
-      default_config_overrides:
-        is_enforced: 0
-        is_active: 0
-  - name: BytesScannedRejectingPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - project_id
-        - referrer
-      default_config_overrides:
-        is_active: 0
-        is_enforced: 0
 
 query_processors:
   - processor: UniqInSelectAndHavingProcessor
diff --git a/snuba/datasets/configuration/events_analytics_platform/storages/eap_items_downsample_512.yaml b/snuba/datasets/configuration/events_analytics_platform/storages/eap_items_downsample_512.yaml
@@ -108,31 +108,6 @@ schema:
   local_table_name: eap_items_1_downsample_512_local
   dist_table_name: eap_items_1_downsample_512_dist
   partition_format: [date]
-allocation_policies:
-  - name: ConcurrentRateLimitAllocationPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - referrer
-        - project_id
-      default_config_overrides:
-        is_enforced: 0
-  - name: ReferrerGuardRailPolicy
-    args:
-      required_tenant_types:
-        - referrer
-      default_config_overrides:
-        is_enforced: 0
-        is_active: 0
-  - name: BytesScannedRejectingPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - project_id
-        - referrer
-      default_config_overrides:
-        is_active: 0
-        is_enforced: 0
 
 query_processors:
   - processor: UniqInSelectAndHavingProcessor
diff --git a/snuba/datasets/configuration/events_analytics_platform/storages/eap_items_downsample_64.yaml b/snuba/datasets/configuration/events_analytics_platform/storages/eap_items_downsample_64.yaml
@@ -109,31 +109,6 @@ schema:
   local_table_name: eap_items_1_downsample_64_local
   dist_table_name: eap_items_1_downsample_64_dist
   partition_format: [date]
-allocation_policies:
-  - name: ConcurrentRateLimitAllocationPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - referrer
-        - project_id
-      default_config_overrides:
-        is_enforced: 0
-  - name: ReferrerGuardRailPolicy
-    args:
-      required_tenant_types:
-        - referrer
-      default_config_overrides:
-        is_enforced: 0
-        is_active: 0
-  - name: BytesScannedRejectingPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - project_id
-        - referrer
-      default_config_overrides:
-        is_active: 0
-        is_enforced: 0
 
 query_processors:
   - processor: UniqInSelectAndHavingProcessor
diff --git a/snuba/datasets/configuration/events_analytics_platform/storages/eap_items_downsample_8.yaml b/snuba/datasets/configuration/events_analytics_platform/storages/eap_items_downsample_8.yaml
@@ -108,31 +108,6 @@ schema:
   local_table_name: eap_items_1_downsample_8_local
   dist_table_name: eap_items_1_downsample_8_dist
   partition_format: [date]
-allocation_policies:
-  - name: ConcurrentRateLimitAllocationPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - referrer
-        - project_id
-      default_config_overrides:
-        is_enforced: 0
-  - name: ReferrerGuardRailPolicy
-    args:
-      required_tenant_types:
-        - referrer
-      default_config_overrides:
-        is_enforced: 0
-        is_active: 0
-  - name: BytesScannedRejectingPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - project_id
-        - referrer
-      default_config_overrides:
-        is_active: 0
-        is_enforced: 0
 
 query_processors:
   - processor: UniqInSelectAndHavingProcessor
diff --git a/snuba/query/allocation_policies/__init__.py b/snuba/query/allocation_policies/__init__.py
@@ -35,6 +35,21 @@
 MAX_THRESHOLD = int(1e12)
 
 
+def get_max_bytes_to_read(quota_allowances: dict[str, QuotaAllowance]) -> int:
+    """
+    alternative implementation:
+    non_zero_max_bytes = [qa.max_bytes_to_read for qa in quota_allowances.values() if qa.max_bytes_to_read > 0]
+    return min(non_zero_max_bytes) if non_zero_max_bytes else 0
+    """
+    max_bytes_to_read = min(
+        [qa.max_bytes_to_read for qa in quota_allowances.values()],
+        key=lambda mb: float("inf") if mb == 0 else mb,
+    )
+    if max_bytes_to_read != 0:
+        return max_bytes_to_read
+    return 0
+
+
 @dataclass(frozen=True)
 class QueryResultOrError:
     """When a query executes, even if it errors, we still want the stats associated
diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py
@@ -32,6 +32,7 @@
     AllocationPolicyViolations,
     QueryResultOrError,
     QuotaAllowance,
+    get_max_bytes_to_read,
 )
 from snuba.query.composite import CompositeQuery
 from snuba.query.data_source.join import IndividualNode, JoinClause, JoinVisitor
@@ -808,7 +809,12 @@ def _apply_allocation_policies_quota(
     Sets the resource quota in the query_settings object to the minimum of all available
     quota allowances from the given allocation policies.
     """
-    quota_allowances: dict[str, Any] = {}
+
+    # if it's an EAP query, then allocation policies were already applied in the routing strategy layer, so we bypass it here
+    if attribution_info.app_id.key == "eap":
+        return
+
+    quota_allowances: dict[str, QuotaAllowance] = {}
     can_run = True
     rejection_quota_and_policy = None
     throttle_quota_and_policy = None
@@ -855,10 +861,7 @@ def _apply_allocation_policies_quota(
         summary: dict[str, Any] = {}
         summary["threads_used"] = min_threads_across_policies
 
-        max_bytes_to_read = min(
-            [qa.max_bytes_to_read for qa in quota_allowances.values()],
-            key=lambda mb: float("inf") if mb == 0 else mb,
-        )
+        max_bytes_to_read = get_max_bytes_to_read(quota_allowances)
         if max_bytes_to_read != 0:
             query_settings.push_clickhouse_setting("max_bytes_to_read", max_bytes_to_read)
             summary["max_bytes_to_read"] = max_bytes_to_read
diff --git a/snuba/web/rpc/__init__.py b/snuba/web/rpc/__init__.py
@@ -31,10 +31,10 @@
     RPCRequestException,
     convert_rpc_exception_to_proto,
 )
-from snuba.web.rpc.storage_routing.load_retriever import get_cluster_loadinfo
 from snuba.web.rpc.storage_routing.routing_strategies.storage_routing import (
     RoutingContext,
     RoutingDecision,
+    get_stats_dict,
 )
 from snuba.web.rpc.storage_routing.routing_strategy_selector import (
     RoutingStrategySelector,
@@ -191,7 +191,9 @@ def execute(self, in_msg: Tin) -> Tout:
         span = scope.span
         if span is not None:
             span.description = self.config_key()
-        self.routing_context = RoutingContext(timer=self._timer, in_msg=in_msg)
+        self.routing_context = RoutingContext(
+            timer=self._timer, in_msg=in_msg, query_id=uuid.uuid4().hex
+        )
 
         self.__before_execute(in_msg)
         error: Exception | None = None
@@ -202,7 +204,29 @@ def execute(self, in_msg: Tin) -> Tout:
                     span.set_data("selected_tier", self.routing_decision.tier)
                     out = self._execute(in_msg)
             else:
-                raise AllocationPolicyViolations
+                self.metrics.increment(
+                    "request_rate_limited",
+                    tags=self._timer.tags,
+                )
+                error = QueryException.from_args(
+                    AllocationPolicyViolations.__name__,
+                    "Query cannot be run due to routing strategy deciding it cannot run, most likely due to allocation policies",
+                    extra={
+                        "stats": get_stats_dict(self.routing_decision),
+                        "sql": "no sql run",
+                        "experiments": {},
+                    },
+                )
+                error.__cause__ = AllocationPolicyViolations.from_args(
+                    {
+                        "summary": {},
+                        "details": {
+                            key: quota_allowance.to_dict()
+                            for key, quota_allowance in self.routing_decision.routing_context.allocation_policies_recommendations.items()
+                        },
+                    }
+                )
+                raise error
         except QueryException as e:
             out = self.response_class()()
             if (
@@ -253,16 +277,10 @@ def __before_execute(self, in_msg: Tin) -> None:
         meta = getattr(in_msg, "meta", None)
         if meta is not None:
             if not hasattr(meta, "request_id") or not meta.request_id:
-                meta.request_id = str(uuid.uuid4())
+                meta.request_id = self.routing_context.query_id
 
         self._timer.update_tags(self.__extract_request_tags(in_msg))
 
-        # we're calling this function to get the cluster load info to emit metrics and to prevent dead code
-        # the result is currently not used in storage routing
-        # can turn off on Snuba Admin
-        if state.get_config("storage_routing.enable_get_cluster_loadinfo", True):
-            get_cluster_loadinfo()
-
         selected_strategy = RoutingStrategySelector().select_routing_strategy(self.routing_context)
         self.routing_decision = selected_strategy.get_routing_decision(self.routing_context)
         self._timer.mark("rpc_start")
@@ -325,6 +343,10 @@ def __after_execute(self, in_msg: Tin, out_msg: Tout, error: Exception | None) -
         output_metrics_error = None
         try:
             self.routing_decision.strategy.output_metrics(self.routing_decision)
+            if self.routing_decision.routing_context.query_result is not None or isinstance(
+                error, QueryException
+            ):
+                self.routing_decision.strategy.update_allocation_policies_balances(self.routing_decision, error)  # type: ignore
         except Exception as e:
             output_metrics_error = e
 
@@ -336,13 +358,11 @@ def __after_execute(self, in_msg: Tin, out_msg: Tout, error: Exception | None) -
                     "request_invalid",
                     tags=self._timer.tags,
                 )
-            elif isinstance(error, AllocationPolicyViolations):
-                sentry_sdk.capture_exception(error)
-                self.metrics.increment(
-                    "request_rate_limited",
-                    tags=self._timer.tags,
-                )
-            else:
+            # AllocationPolicyViolations is not a request_error
+            elif not (
+                isinstance(error, QueryException)
+                and error.exception_type == AllocationPolicyViolations.__name__
+            ):
                 sentry_sdk.capture_exception(error)
                 self.metrics.increment(
                     "request_error",
diff --git a/snuba/web/rpc/storage_routing/routing_strategies/outcomes_based.py b/snuba/web/rpc/storage_routing/routing_strategies/outcomes_based.py
@@ -149,14 +149,26 @@ def _get_min_timerange_to_query_outcomes(self) -> int:
             or default
         )
 
-    def _get_routing_decision(self, routing_context: RoutingContext) -> RoutingDecision:
+    def _get_routing_decision(
+        self,
+        routing_context: RoutingContext,
+    ) -> RoutingDecision:
+        combined_allocation_policies_recommendations = (
+            self._get_combined_allocation_policies_recommendations(routing_context)
+        )
+
         routing_decision = RoutingDecision(
             routing_context=routing_context,
             strategy=self,
             tier=Tier.TIER_1,
-            clickhouse_settings={},
-            can_run=True,
+            clickhouse_settings=combined_allocation_policies_recommendations["settings"],
+            can_run=combined_allocation_policies_recommendations["can_run"],
+            is_throttled=combined_allocation_policies_recommendations["is_throttled"],
         )
+
+        if not routing_decision.can_run:
+            return routing_decision
+
         in_msg_meta = extract_message_meta(routing_decision.routing_context.in_msg)
         sentry_sdk.update_current_span(
             attributes={
diff --git a/snuba/web/rpc/storage_routing/routing_strategies/storage_routing.py b/snuba/web/rpc/storage_routing/routing_strategies/storage_routing.py
diff --git a/tests/web/rpc/v1/routing_strategies/test_outcomes_based.py b/tests/web/rpc/v1/routing_strategies/test_outcomes_based.py
diff --git a/tests/web/rpc/v1/routing_strategies/test_strategy_selector.py b/tests/web/rpc/v1/routing_strategies/test_strategy_selector.py
diff --git a/tests/web/rpc/v1/test_storage_routing.py b/tests/web/rpc/v1/test_storage_routing.py
