def test_delegate( settings_class: Type[Union[HTTPQuerySettings, SubscriptionQuerySettings]], expected_rate_limiters: Sequence[RateLimitParameters], ) -> None: settings = settings_class( referrer="test", consistent=False, ) settings.add_rate_limit( RateLimitParameters( rate_limit_name="rate_name", bucket="project", per_second_limit=10.0, concurrent_limit=22, ) ) settings_delegate = RateLimiterDelegate("secondary", settings) settings_delegate.add_rate_limit( RateLimitParameters( rate_limit_name="second_rate_name", bucket="table", per_second_limit=11.0, concurrent_limit=23, ) ) assert settings_delegate.referrer == settings.referrer assert settings_delegate.get_rate_limit_params() == expected_rate_limiters
def test_per_second_limit(self): bucket = uuid.uuid4() rate_limit_params = RateLimitParameters("foo", bucket, 1, None) # Create 30 queries at time 0, should all be allowed with patch.object(state.time, "time", lambda: 0): for _ in range(30): with rate_limit(rate_limit_params) as stats: assert stats is not None # Create another 30 queries at time 30, should also be allowed with patch.object(state.time, "time", lambda: 30): for _ in range(30): with rate_limit(rate_limit_params) as stats: assert stats is not None with patch.object(state.time, "time", lambda: 60): # 1 more query should be allowed at T60 because it does not make the previous # rate exceed 1/sec until it has finished. with rate_limit(rate_limit_params) as stats: assert stats is not None # But the next one should not be allowed with pytest.raises(RateLimitExceeded): with rate_limit(rate_limit_params): pass # Another query at time 61 should be allowed because the first 30 queries # have fallen out of the lookback window with patch.object(state.time, "time", lambda: 61): with rate_limit(rate_limit_params) as stats: assert stats is not None
def process_query(self, query: Query, request_settings: RequestSettings) -> None: # If the settings don't already have a project rate limit, add one existing = request_settings.get_rate_limit_params() for ex in existing: if ex.rate_limit_name == PROJECT_RATE_LIMIT_NAME: return project_ids = get_project_ids_in_query_ast(query, self.project_column) if not project_ids: return # TODO: Use all the projects, not just one project_id = project_ids.pop() prl, pcl = get_configs([("project_per_second_limit", 1000), ("project_concurrent_limit", 1000)]) # Specific projects can have their rate limits overridden (per_second, concurr) = get_configs([ ("project_per_second_limit_{}".format(project_id), prl), ("project_concurrent_limit_{}".format(project_id), pcl), ]) rate_limit = RateLimitParameters( rate_limit_name=PROJECT_RATE_LIMIT_NAME, bucket=str(project_id), per_second_limit=per_second, concurrent_limit=concurr, ) request_settings.add_rate_limit(rate_limit)
def process_query(self, query: Query, query_settings: QuerySettings) -> None: # If the settings don't already have an object rate limit, add one if self._is_already_applied(query_settings): return per_second_name = self.get_per_second_name(query, query_settings) concurrent_name = self.get_concurrent_name(query, query_settings) object_rate_limit, object_concurrent_limit = get_configs( [ (per_second_name, self.default_limit), (concurrent_name, self.default_limit), ] ) obj_id = self.get_object_id(query, query_settings) if obj_id is None: return # Specific objects can have their rate limits overridden (per_second, concurr) = get_configs( [ (f"{per_second_name}_{obj_id}", object_rate_limit), (f"{concurrent_name}_{obj_id}", object_concurrent_limit), ] ) rate_limit = RateLimitParameters( rate_limit_name=self.rate_limit_name, bucket=str(obj_id), per_second_limit=per_second, concurrent_limit=concurr, ) query_settings.add_rate_limit(rate_limit)
def __append_prefix( self, rate_limiter: RateLimitParameters) -> RateLimitParameters: return RateLimitParameters( rate_limit_name=rate_limiter.rate_limit_name, bucket=f"{self.__prefix}_{rate_limiter.bucket}", per_second_limit=rate_limiter.per_second_limit, concurrent_limit=rate_limiter.concurrent_limit, )
def test_aggregator(self): # do not raise with multiple valid rate limits rate_limit_params_outer = RateLimitParameters("foo", "bar", None, 5) rate_limit_params_inner = RateLimitParameters("foo", "bar", None, 5) with RateLimitAggregator( [rate_limit_params_outer, rate_limit_params_inner]): pass # raise when the inner rate limit should fail rate_limit_params_outer = RateLimitParameters("foo", "bar", None, 0) rate_limit_params_inner = RateLimitParameters("foo", "bar", None, 5) with pytest.raises(RateLimitExceeded): with RateLimitAggregator( [rate_limit_params_outer, rate_limit_params_inner]): pass # raise when the outer rate limit should fail rate_limit_params_outer = RateLimitParameters("foo", "bar", None, 5) rate_limit_params_inner = RateLimitParameters("foo", "bar", None, 0) with pytest.raises(RateLimitExceeded): with RateLimitAggregator( [rate_limit_params_outer, rate_limit_params_inner]): pass
def process_query(self, query: Query, query_settings: QuerySettings) -> None: table_name = query.get_from_clause().table_name (per_second, concurr) = get_configs([ (f"table_per_second_limit_{table_name}{self.__suffix}", 5000), (f"table_concurrent_limit_{table_name}{self.__suffix}", 1000), ]) rate_limit = RateLimitParameters( rate_limit_name=TABLE_RATE_LIMIT_NAME, bucket=table_name, per_second_limit=per_second, concurrent_limit=concurr, ) query_settings.add_rate_limit(rate_limit)
def _get_rate_limit_params( self, project_ids: Sequence[int]) -> RateLimitParameters: project_id = (project_ids[0] if project_ids else 0 ) # TODO rate limit on every project in the list? prl, pcl = get_configs([("project_per_second_limit", 1000), ("project_concurrent_limit", 1000)]) # Specific projects can have their rate limits overridden (per_second, concurr) = get_configs([ ("project_per_second_limit_{}".format(project_id), prl), ("project_concurrent_limit_{}".format(project_id), pcl), ]) return RateLimitParameters( rate_limit_name=PROJECT_RATE_LIMIT_NAME, bucket=str(project_id), per_second_limit=per_second, concurrent_limit=concurr, )
def test_concurrent_limit(self): # No concurrent limit should not raise rate_limit_params = RateLimitParameters("foo", "bar", None, None) with rate_limit(rate_limit_params) as stats: assert stats is not None # 0 concurrent limit rate_limit_params = RateLimitParameters("foo", "bar", None, 0) with pytest.raises(RateLimitExceeded): with rate_limit(rate_limit_params): pass # Concurrent limit 1 with consecutive queries should not raise rate_limit_params = RateLimitParameters("foo", "bar", None, 1) with rate_limit(rate_limit_params): pass with rate_limit(rate_limit_params): pass # Concurrent limit with concurrent queries rate_limit_params = RateLimitParameters("foo", "bar", None, 1) with pytest.raises(RateLimitExceeded): with rate_limit(rate_limit_params): with rate_limit(rate_limit_params): pass # Concurrent with different buckets should not raise rate_limit_params1 = RateLimitParameters("foo", "bar", None, 1) rate_limit_params2 = RateLimitParameters("shoe", "star", None, 1) with RateLimitAggregator([rate_limit_params1]): with RateLimitAggregator([rate_limit_params2]): pass
def test_concurrent_limit(self): # No concurrent limit should not raise rate_limit_params = RateLimitParameters('foo', 'bar', None, None) with rate_limit(rate_limit_params): pass # 0 concurrent limit rate_limit_params = RateLimitParameters('foo', 'bar', None, 0) with pytest.raises(RateLimitExceeded): with rate_limit(rate_limit_params): pass # Concurrent limit 1 with consecutive queries should not raise rate_limit_params = RateLimitParameters('foo', 'bar', None, 1) with rate_limit(rate_limit_params): pass with rate_limit(rate_limit_params): pass # Concurrent limit with concurrent queries rate_limit_params = RateLimitParameters('foo', 'bar', None, 1) with pytest.raises(RateLimitExceeded): with rate_limit(rate_limit_params): with rate_limit(rate_limit_params): pass # Concurrent with different buckets should not raise rate_limit_params1 = RateLimitParameters('foo', 'bar', None, 1) rate_limit_params2 = RateLimitParameters('shoe', 'star', None, 1) with RateLimitAggregator([rate_limit_params1]): with RateLimitAggregator([rate_limit_params2]): pass
def test_bypass_rate_limit(self): rate_limit_params = RateLimitParameters("foo", "bar", None, None) state.set_config("bypass_rate_limit", 1) with rate_limit(rate_limit_params) as stats: assert stats is None
from typing import Sequence, Type, Union import pytest from snuba.pipeline.settings_delegator import RateLimiterDelegate from snuba.query.query_settings import HTTPQuerySettings, SubscriptionQuerySettings from snuba.state.rate_limit import RateLimitParameters test_cases = [ pytest.param( HTTPQuerySettings, [ RateLimitParameters( rate_limit_name="rate_name", bucket="secondary_project", per_second_limit=10.0, concurrent_limit=22, ), RateLimitParameters( rate_limit_name="second_rate_name", bucket="secondary_table", per_second_limit=11.0, concurrent_limit=23, ), ], id="HTTP Request Settings", ), pytest.param( SubscriptionQuerySettings, [], id="Subscriptions request.query_settings" ), ]
from snuba.query.data_source.simple import Table from snuba.query.processors.table_rate_limit import TableRateLimit from snuba.query.query_settings import HTTPQuerySettings from snuba.state import set_config from snuba.state.rate_limit import TABLE_RATE_LIMIT_NAME, RateLimitParameters test_data = [ pytest.param( TableRateLimit(), Query(Table("errors_local", ColumnSet([])), selected_columns=[], condition=None), "table_concurrent_limit_transactions_local", RateLimitParameters( rate_limit_name=TABLE_RATE_LIMIT_NAME, bucket="errors_local", per_second_limit=5000, concurrent_limit=1000, ), id="Set rate limiter on another table", ), pytest.param( TableRateLimit(), Query(Table("errors_local", ColumnSet([])), selected_columns=[], condition=None), "table_concurrent_limit_errors_local", RateLimitParameters( rate_limit_name=TABLE_RATE_LIMIT_NAME, bucket="errors_local", per_second_limit=5000, concurrent_limit=50,