diff --git a/rate_limit/provider.py b/rate_limit/provider.py index 46e9d2c..d7da105 100644 --- a/rate_limit/provider.py +++ b/rate_limit/provider.py @@ -68,6 +68,8 @@ def __init__(self, service_type, logger=log.Logger(__name__), **kwargs): super(ConfigurationRateLimitProvider, self).__init__( service_type=service_type, logger=logger, kwargs=kwargs ) + # Custom rate limits per scope. + self.custom_ratelimits = {} def get_global_rate_limits(self, action, target_type_uri, **kwargs): """ @@ -101,7 +103,10 @@ def get_local_rate_limits(self, scope, action, target_type_uri, **kwargs): :param kwargs: optional, additional parameters :return: the local rate limit or -1 if not set """ - ttu_ratelimits = self.local_ratelimits.get(target_type_uri, []) + _ratelimits = self.custom_ratelimits.get(scope, None) + if not _ratelimits: + _ratelimits = self.local_ratelimits + ttu_ratelimits = _ratelimits.get(target_type_uri, []) if not ttu_ratelimits: ttu_ratelimits = self._get_wildcard_ratelimits( self.local_ratelimits, @@ -157,6 +162,9 @@ def read_rate_limits_from_config(self, config_path): rates = config.get('rates', {}) self.global_ratelimits = rates.get('global', {}) self.local_ratelimits = rates.get('default', {}) + for scope, rl in rates.items(): + if scope not in ['global', 'default']: + self.custom_ratelimits[scope] = rl class LimesRateLimitProvider(RateLimitProvider):