class SlidingRedisBackend(BaseBackend):
def __init__(
self,
host: str = "localhost",
port: int = 6379,
db: int = 0,
password: str = None,
) -> None:
self._redis = StrictRedis(host=host,
port=port,
db=db,
password=password)
self.sliding_function = self._redis.register_script(
SLIDING_WINDOW_SCRIPT)
async def get_limits(self, path: str, user: str, rule: Rule) -> bool:
epoch = time.time()
ruleset = rule.ruleset(path, user)
keys = list(ruleset.keys())
args = [epoch, json.dumps(ruleset)]
# quoted_args = [f"'{a}'" for a in args]
# cli = f"redis-cli --ldb --eval /tmp/script.lua {' '.join(keys)} , {' '.join(quoted_args)}"
# logger.debug(cli)
r = await self.sliding_function.execute(keys=keys, args=args)
# from tests.backends.test_redis import logger
# logger.debug(f"{epoch} {r} : {all(r)}")
return all(r)
async def decrease_limit(self, path: str, user: str, rule: Rule) -> bool:
raise NotImplementedError()
async def increase_limit(self, path: str, user: str, rule: Rule) -> bool:
raise NotImplementedError()
async def set_block_time(self, user: str, block_time: int) -> None:
await self._redis.set(f"blocking:{user}", True, block_time)
async def is_blocking(self, user: str) -> bool:
return bool(await self._redis.get(f"blocking:{user}"))
async def allow_request(self, path: str, user: str, rule: Rule) -> bool:
if await self.is_blocking(user):
return False
allow = await self.get_limits(path, user, rule)
if not allow and rule.block_time:
await self.set_block_time(user, rule.block_time)
return allow
class RedisBackend(BaseBackend):
def __init__(
self,
host: str = "localhost",
port: int = 6379,
db: int = 0,
password: str = None,
) -> None:
self._redis = StrictRedis(host=host, port=port, db=db, password=password)
self.decrease_function = self._redis.register_script(DECREASE_SCRIPT)
async def increase_limit(self, path: str, user: str, rule: Rule) -> bool:
"""
Return True means successful increase.
"""
async with await self._redis.pipeline() as pipe: # type: StrictPipeline
try:
timestamp = time.time()
await pipe.watch(
*[f"{path}:{user}:{name}:last_modify" for name in RULENAMES]
)
pipe.multi()
[
await pipe.get(f"{path}:{user}:{name}:last_modify")
for name in RULENAMES
]
result = [
None if _timestamp is None else float(_timestamp)
for _timestamp in await pipe.execute()
]
incr_dict = self.calc_incr_value(result, rule)
if not incr_dict:
return False
pipe.multi()
for name, data in incr_dict.items():
await pipe.set(f"{path}:{user}:{name}", data["value"], data["ttl"])
await pipe.set(
f"{path}:{user}:{name}:last_modify", timestamp, data["ttl"],
)
await pipe.execute()
return True
except WatchError: # pragma: no cover
return False
async def decrease_limit(self, path: str, user: str, rule: Rule) -> bool:
"""
Return True means successful decrease.
"""
return await self.decrease_function.execute(
keys=[
f"{path}:{user}:{name}"
for name in RULENAMES
if getattr(rule, name) is not None
]
)
async def set_block_time(self, user: str, block_time: int) -> None:
await self._redis.set(f"blocking:{user}", True, block_time)
async def is_blocking(self, user: str) -> bool:
return bool(await self._redis.get(f"blocking:{user}"))