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,