ratelimit: comptabilise les messages WebSocket (stream.accounting) dans le quota de lignes de la clé
1 changed file +42 −0
modified
hfmarketdata/api/ratelimit/middleware.py
+42 −0
@@ -371,8 +371,50 @@ class RateLimitMiddleware: | ||
| 371 | 371 | return await self.app(scope, receive, send) |
| 372 | 372 | |
| 373 | 373 | |
| 374 | +def _tier_for_principal(principal: str): | |
| 375 | + """Tier of a `key:<id>` principal (SQLite lookup, 60 s cache) — keyless principals get the keyless tier.""" | |
| 376 | + from . import tiers as _tiers | |
| 377 | + if not principal.startswith("key:"): | |
| 378 | + return _tiers.TIERS["keyless"] | |
| 379 | + now = time.time() | |
| 380 | + with _key_cache_lock: | |
| 381 | + hit = _key_cache.get(principal) | |
| 382 | + if hit and hit[0] > now: | |
| 383 | + return hit[1] | |
| 384 | + tier_name = "free" | |
| 385 | + try: | |
| 386 | + from sqlalchemy import select | |
| 387 | + from accounts.models import ApiKey, User | |
| 388 | + from core.db import session | |
| 389 | + with session() as s: | |
| 390 | + row = s.execute(select(ApiKey.tier_override, User.tier).join(User, User.id == ApiKey.user_id) | |
| 391 | + .where(ApiKey.id == int(principal[4:]))).first() | |
| 392 | + if row: | |
| 393 | + tier_name = row[0] or row[1] or "free" | |
| 394 | + except Exception: | |
| 395 | + pass | |
| 396 | + tier = _tiers.TIERS.get(tier_name, _tiers.TIERS["free"]) | |
| 397 | + with _key_cache_lock: | |
| 398 | + _key_cache[principal] = (now + KEY_CACHE_TTL_S, tier) | |
| 399 | + return tier | |
| 400 | + | |
| 401 | + | |
| 402 | +def stream_charger(principal: str, rows: int) -> None: | |
| 403 | + """Charge WebSocket messages to the rows quota of the key (1 row per message, no request charge).""" | |
| 404 | + from . import redis_limiter | |
| 405 | + try: | |
| 406 | + redis_limiter.apply_tier(principal, _tier_for_principal(principal), req_cost=0, rows_cost=int(rows), force=True) | |
| 407 | + except Exception: # never let accounting break a live stream | |
| 408 | + pass | |
| 409 | + | |
| 410 | + | |
| 374 | 411 | def install(app: FastAPI) -> None: |
| 375 | 412 | """Mount the middleware INSIDE the CORS layer so 429/401 responses keep their CORS headers.""" |
| 413 | + try: # stream module (fundamentals filings WebSocket) → charge delivered messages to the rows quota | |
| 414 | + from stream import accounting as _stream_accounting | |
| 415 | + _stream_accounting.set_charger(stream_charger) | |
| 416 | + except Exception: | |
| 417 | + pass | |
| 376 | 418 | if app.middleware_stack is not None: # pragma: no cover — must run before the first request |
| 377 | 419 | raise RuntimeError("ratelimit.middleware.install() must be called before the app starts") |
| 378 | 420 | app.user_middleware.append(Middleware(RateLimitMiddleware)) |
| 379 | 421 | |