diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py index 0fc9acebf2..a8281cfe85 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -914,16 +914,32 @@ def _apply_allocation_policies_quota( stats["quota_allowance"]["summary"] = summary if not can_run: + rejecting_policy = ( + rejection_quota_and_policy.policy.class_name() + if rejection_quota_and_policy is not None + else "unknown" + ) + span.set_data("policy", rejecting_policy) + span.set_data("action", "rejected") metrics.increment( "rejected_query", - tags={"storage_key": allocation_policies[0].resource_identifier.value}, + tags={ + "storage_key": allocation_policies[0].resource_identifier.value, + "policy": rejecting_policy, + }, ) raise AllocationPolicyViolations.from_args(stats["quota_allowance"]) if throttle_quota_and_policy is not None: + throttling_policy = throttle_quota_and_policy.policy.class_name() + span.set_data("policy", throttling_policy) + span.set_data("action", "throttled") metrics.increment( "throttled_query", - tags={"storage_key": allocation_policies[0].resource_identifier.value}, + tags={ + "storage_key": allocation_policies[0].resource_identifier.value, + "policy": throttling_policy, + }, ) else: metrics.increment( diff --git a/tests/web/test_db_query.py b/tests/web/test_db_query.py index eb1db761d2..ffce03fe2b 100644 --- a/tests/web/test_db_query.py +++ b/tests/web/test_db_query.py @@ -827,6 +827,12 @@ def __init__(self, max_threads: int, policy_name: str) -> None: }, } } + throttled_metrics = get_recorded_metric_calls("increment", "db_query.throttled_query") + assert throttled_metrics + assert throttled_metrics[0].tags == { + "storage_key": "doesntmatter", + "policy": "ThrottleAllocationPolicy1", + } def test_db_query_with_rejecting_allocation_policy() -> None: @@ -950,6 +956,12 @@ def _update_quota_balance( assert update_called, ( "update_quota_balance should have been called even though the query was rejected but was not" ) + rejected_metrics = get_recorded_metric_calls("increment", "db_query.rejected_query") + assert rejected_metrics + assert rejected_metrics[0].tags == { + "storage_key": "doesntmatter", + "policy": "RejectAllocationPolicy", + } @pytest.mark.events_db