59 lines
2.3 KiB
Python
59 lines
2.3 KiB
Python
|
"""NL2SQL 查询配额、并发和频率限制。"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import time
|
||
|
|
|
||
|
|
|
||
|
|
class QueryLimiter:
|
||
|
|
"""基于 Redis 计数器的请求限制器，Redis 故障时允许降级放行。"""
|
||
|
|
|
||
|
|
def __init__(self, redis, *, rate_window_seconds: int = 60):
|
||
|
|
self.redis = redis
|
||
|
|
self.rate_window_seconds = rate_window_seconds
|
||
|
|
|
||
|
|
async def acquire(
|
||
|
|
self,
|
||
|
|
user_id: int,
|
||
|
|
*,
|
||
|
|
daily_quota: int = 0,
|
||
|
|
max_concurrent: int = 1,
|
||
|
|
rate_limit: int = 0,
|
||
|
|
) -> bool:
|
||
|
|
"""检查并占用一次配额、并发和频率额度，0 表示不限制。"""
|
||
|
|
try:
|
||
|
|
if daily_quota:
|
||
|
|
daily_key = f"nl2sql:daily:{user_id}:{time.strftime('%Y%m%d')}"
|
||
|
|
daily_count = await self.redis.incr(daily_key)
|
||
|
|
await self.redis.expire(daily_key, 86400)
|
||
|
|
if daily_count > daily_quota:
|
||
|
|
await self.redis.decr(daily_key)
|
||
|
|
return False
|
||
|
|
concurrent_key = f"nl2sql:concurrent:{user_id}"
|
||
|
|
concurrent_count = await self.redis.incr(concurrent_key)
|
||
|
|
await self.redis.expire(concurrent_key, 3600)
|
||
|
|
if max_concurrent and concurrent_count > max_concurrent:
|
||
|
|
await self.redis.decr(concurrent_key)
|
||
|
|
if daily_quota:
|
||
|
|
await self.redis.decr(daily_key)
|
||
|
|
return False
|
||
|
|
if rate_limit:
|
||
|
|
rate_key = f"nl2sql:rate:{user_id}:{int(time.time()) // self.rate_window_seconds}"
|
||
|
|
rate_count = await self.redis.incr(rate_key)
|
||
|
|
await self.redis.expire(rate_key, self.rate_window_seconds + 1)
|
||
|
|
if rate_count > rate_limit:
|
||
|
|
await self.redis.decr(rate_key)
|
||
|
|
await self.redis.decr(concurrent_key)
|
||
|
|
if daily_quota:
|
||
|
|
await self.redis.decr(daily_key)
|
||
|
|
return False
|
||
|
|
return True
|
||
|
|
except Exception: # noqa: BLE001 Redis 故障时降级为不限制
|
||
|
|
return True
|
||
|
|
|
||
|
|
async def release(self, user_id: int) -> None:
|
||
|
|
"""释放一次并发额度，失败时静默降级。"""
|
||
|
|
try:
|
||
|
|
await self.redis.decr(f"nl2sql:concurrent:{user_id}")
|
||
|
|
except Exception: # noqa: BLE001 Redis 故障不能影响查询结果
|
||
|
|
return
|