SPB Git

spb/findata

Public
1.4 KB · 35 lines python
1"""In-process token bucket per principal (single API process). Quotas (rows/day) live in usage_daily."""2from __future__ import annotations34import threading5import time67from findata_core.db import transaction8910class TokenBucket:11    def __init__(self):12        self._b: dict[str, tuple[float, float]] = {}13        self._lock = threading.Lock()1415    def allow(self, key: str, per_minute: int) -> tuple[bool, int]:16        now = time.monotonic()17        with self._lock:18            tokens, last = self._b.get(key, (float(per_minute), now))19            tokens = min(per_minute, tokens + (now - last) * per_minute / 60.0)20            if tokens >= 1:21                self._b[key] = (tokens - 1, now)22                return True, int(tokens - 1)23            self._b[key] = (tokens, now)24            return False, 0252627bucket = TokenBucket()282930def record_usage(principal: str, rows: int = 0, bytes_exported: int = 0) -> None:31    with transaction() as conn:32        conn.execute("INSERT INTO usage_daily(principal, day, requests, rows_returned, bytes_exported) VALUES (%s, current_date, 1, %s, %s) "33                     "ON CONFLICT (principal, day) DO UPDATE SET requests=usage_daily.requests+1, rows_returned=usage_daily.rows_returned+EXCLUDED.rows_returned, "34                     "bytes_exported=usage_daily.bytes_exported+EXCLUDED.bytes_exported", (principal, rows, bytes_exported))35
Files · 0467d449e66a5e62fb10ba4fcbf2f09fc3d1ef38
Menu

Theme