""" Usage Tracker for ARF API – quotas, tiers, and audit logging. Thread‑safe, atomic quota consumption, idempotent, fail‑closed. """ import json import sqlite3 import threading import time from contextlib import contextmanager from datetime import datetime, timedelta from dataclasses import dataclass from typing import Dict, Any, Optional, List, Tuple from enum import Enum from fastapi import BackgroundTasks, HTTPException, Request # Optional Redis support try: import redis REDIS_AVAILABLE = True except ImportError: REDIS_AVAILABLE = False redis = None class Tier(str, Enum): FREE = "free" PRO = "pro" PREMIUM = "premium" ENTERPRISE = "enterprise" @property def monthly_evaluation_limit(self) -> Optional[int]: limits = { Tier.FREE: 1000, Tier.PRO: 10_000, Tier.PREMIUM: 50_000, Tier.ENTERPRISE: None, # unlimited } return limits[self] @property def audit_log_retention_days(self) -> int: retention = { Tier.FREE: 7, Tier.PRO: 30, Tier.PREMIUM: 90, Tier.ENTERPRISE: 365, } return retention[self] @dataclass class UsageRecord: """Single evaluation usage record.""" api_key: str tier: Tier timestamp: float endpoint: str request_body: Optional[Dict[str, Any]] = None response: Optional[Dict[str, Any]] = None error: Optional[str] = None processing_ms: Optional[float] = None class UsageTracker: """ Thread‑safe usage tracker with atomic quota consumption and idempotency. """ def __init__(self, db_path: str = "arf_usage.db", redis_url: Optional[str] = None): self.db_path = db_path self._local = threading.local() self._init_db() self._redis_client = None if redis_url and REDIS_AVAILABLE: self._redis_client = redis.from_url(redis_url) elif redis_url: raise ImportError( "Redis client not installed. Run: pip install redis") @contextmanager def _get_conn(self): """Get a thread‑local SQLite connection with write‑ahead logging and immediate transactions.""" if not hasattr(self._local, "conn"): self._local.conn = sqlite3.connect( self.db_path, check_same_thread=False, isolation_level=None) self._local.conn.row_factory = sqlite3.Row self._local.conn.execute("PRAGMA journal_mode=WAL") yield self._local.conn def _init_db(self): with self._get_conn() as conn: conn.execute(""" CREATE TABLE IF NOT EXISTS api_keys ( key TEXT PRIMARY KEY, tier TEXT NOT NULL, created_at REAL NOT NULL, last_used_at REAL, is_active INTEGER DEFAULT 1 ) """) conn.execute(""" CREATE TABLE IF NOT EXISTS usage_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, api_key TEXT NOT NULL, tier TEXT NOT NULL, timestamp REAL NOT NULL, endpoint TEXT NOT NULL, request_body TEXT, response TEXT, error TEXT, processing_ms REAL, idempotency_key TEXT UNIQUE ) """) conn.execute(""" CREATE INDEX IF NOT EXISTS idx_api_key_timestamp ON usage_log (api_key, timestamp) """) conn.execute(""" CREATE TABLE IF NOT EXISTS monthly_counts ( api_key TEXT NOT NULL, year_month TEXT NOT NULL, count INTEGER DEFAULT 0, PRIMARY KEY (api_key, year_month) ) """) conn.execute(""" CREATE TABLE IF NOT EXISTS idempotency_keys ( key TEXT PRIMARY KEY, consumed_at REAL NOT NULL ) """) conn.commit() def _get_month_key(self) -> str: return datetime.now().strftime("%Y-%m") def get_or_create_api_key(self, key: str, tier: Tier = Tier.FREE) -> bool: """Register a new API key. Returns True if key exists or was created.""" with self._get_conn() as conn: row = conn.execute( "SELECT key FROM api_keys WHERE key = ?", (key,)).fetchone() if row: return True conn.execute( "INSERT INTO api_keys (key, tier, created_at, is_active) VALUES (?, ?, ?, ?)", (key, tier.value, time.time(), 1) ) conn.commit() return True def get_tier(self, api_key: str) -> Optional[Tier]: """Return the tier for a given API key, or None if key invalid/inactive.""" with self._get_conn() as conn: row = conn.execute( "SELECT tier FROM api_keys WHERE key = ? AND is_active = 1", (api_key,) ).fetchone() if not row: return None return Tier(row["tier"]) def update_api_key_tier(self, api_key: str, new_tier: Tier) -> bool: """Update the tier of an existing API key. Returns True if successful.""" with self._get_conn() as conn: row = conn.execute( "SELECT key FROM api_keys WHERE key = ?", (api_key,)).fetchone() if not row: return False conn.execute( "UPDATE api_keys SET tier = ? WHERE key = ?", (new_tier.value, api_key)) conn.commit() return True # -------------------------------------------------------------------------- # Atomic quota consumption # -------------------------------------------------------------------------- def _consume_quota_atomic_sqlite( self, api_key: str, tier: Tier, month: str) -> bool: # noqa: E501 """ Atomically increment counter only if under limit. Returns True if quota was consumed, False if limit reached. """ limit = tier.monthly_evaluation_limit if limit is None: # Unlimited – still increment for tracking but always succeed with self._get_conn() as conn: conn.execute( """INSERT INTO monthly_counts (api_key, year_month, count) VALUES (?, ?, 1) ON CONFLICT(api_key, year_month) DO UPDATE SET count = count + 1""", (api_key, month) ) conn.commit() return True # Use BEGIN IMMEDIATE to lock the database for the transaction with self._get_conn() as conn: conn.execute("BEGIN IMMEDIATE") try: # Get current count (or 0) row = conn.execute( "SELECT count FROM monthly_counts WHERE api_key = ? AND year_month = ?", (api_key, month) ).fetchone() current = row["count"] if row else 0 if current >= limit: conn.rollback() return False # Increment conn.execute( """INSERT INTO monthly_counts (api_key, year_month, count) VALUES (?, ?, 1) ON CONFLICT(api_key, year_month) DO UPDATE SET count = count + 1""", (api_key, month) ) conn.commit() return True except Exception: conn.rollback() raise def _consume_quota_atomic_redis( self, api_key: str, tier: Tier, month: str) -> bool: """Atomic Lua script for Redis: INCR only if below limit.""" limit = tier.monthly_evaluation_limit if limit is None: # Unlimited – just increment and return True redis_key = f"arf:quota:{api_key}:{month}" self._redis_client.incr(redis_key) self._redis_client.expire(redis_key, timedelta(days=31)) return True lua_script = """ local key = KEYS[1] local limit = tonumber(ARGV[1]) local current = redis.call('GET', key) if current and tonumber(current) >= limit then return 0 end local new = redis.call('INCR', key) redis.call('EXPIRE', key, 2678400) -- 31 days return 1 """ redis_key = f"arf:quota:{api_key}:{month}" result = self._redis_client.eval(lua_script, 1, redis_key, limit) return result == 1 # -------------------------------------------------------------------------- # Idempotency handling # -------------------------------------------------------------------------- def _is_idempotent_key_used(self, key: str) -> bool: """Check if idempotency key already processed.""" with self._get_conn() as conn: row = conn.execute( "SELECT 1 FROM idempotency_keys WHERE key = ?", (key,)).fetchone() return row is not None def _mark_idempotent_key_used(self, key: str, ttl_seconds: int = 86400): """Store idempotency key with expiration (cleanup later).""" with self._get_conn() as conn: conn.execute( "INSERT INTO idempotency_keys (key, consumed_at) VALUES (?, ?)", (key, time.time()) ) conn.commit() # Optionally schedule cleanup of old keys (can be done in a background # thread) # -------------------------------------------------------------------------- # Core usage recording (atomic + idempotent) # -------------------------------------------------------------------------- def consume_quota_and_log( self, record: UsageRecord, idempotency_key: Optional[str] = None, ) -> Tuple[bool, Optional[Dict[str, Any]]]: """ Atomically consume quota and insert audit log. Returns (success, existing_response) where existing_response is not None only when idempotency_key matched a previous successful call. """ # Idempotency check (if key provided) if idempotency_key: if self._is_idempotent_key_used(idempotency_key): # Retrieve previous response from audit log (simplified – you may cache full response) # For full idempotency, we would store the response body in idempotency table. # Here we return a marker that caller should use cached # response. return False, {"idempotent": True, "message": "Already processed"} month = self._get_month_key() # Atomic quota consumption if self._redis_client: quota_ok = self._consume_quota_atomic_redis( record.api_key, record.tier, month) else: quota_ok = self._consume_quota_atomic_sqlite( record.api_key, record.tier, month) if not quota_ok: return False, None # Insert audit log (with idempotency key as unique constraint) try: with self._get_conn() as conn: conn.execute( """INSERT INTO usage_log (api_key, tier, timestamp, endpoint, request_body, response, error, processing_ms, idempotency_key) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", (record.api_key, record.tier.value, record.timestamp, record.endpoint, json.dumps( record.request_body) if record.request_body else None, json.dumps( record.response) if record.response else None, record.error, record.processing_ms, idempotency_key, )) conn.commit() except sqlite3.IntegrityError as e: # Duplicate idempotency_key – already inserted by another # concurrent request if "UNIQUE constraint failed: usage_log.idempotency_key" in str(e): return False, {"idempotent": True, "message": "Already processed"} raise if idempotency_key: self._mark_idempotent_key_used(idempotency_key) # Removed stray # noqa: E501 comment that was wrongly indented here return True, None # -------------------------------------------------------------------------- # Legacy interface (kept for compatibility but deprecated) # -------------------------------------------------------------------------- def increment_usage_sync( self, record: UsageRecord, idempotency_key: Optional[str] = None) -> bool: """ Synchronously record usage and increment counter. Returns True if within quota and recorded, False otherwise. This method now uses the atomic implementation. """ success, _ = self.consume_quota_and_log(record, idempotency_key) return success async def increment_usage_async( self, record: UsageRecord, background_tasks: BackgroundTasks, idempotency_key: Optional[str] = None ) -> bool: """ Asynchronously record usage using FastAPI BackgroundTasks. Still does the atomic check synchronously, then schedules the insert. """ # First, do atomic quota check (synchronous) – we must ensure we don't double-consume. # Because background tasks may run later, we still need to reserve quota now. # Simplified: we call consume_quota_and_log synchronously – that defeats async benefit. # Better to use a queue or Redis with background processing. # For this fix, we'll use the sync method (blocking) but still support # idempotency. return self.increment_usage_sync(record, idempotency_key) # -------------------------------------------------------------------------- # Quota inspection (non‑atomic, for display only) # -------------------------------------------------------------------------- def get_remaining_quota(self, api_key: str, tier: Tier) -> Optional[int]: """Return remaining evaluations for the month (non‑atomic, for info only).""" limit = tier.monthly_evaluation_limit if limit is None: return None month = self._get_month_key() if self._redis_client: redis_key = f"arf:quota:{api_key}:{month}" count = int(self._redis_client.get(redis_key) or 0) return max(0, limit - count) with self._get_conn() as conn: row = conn.execute( "SELECT count FROM monthly_counts WHERE api_key = ? AND year_month = ?", (api_key, month) ).fetchone() count = row["count"] if row else 0 return max(0, limit - count) # -------------------------------------------------------------------------- # Audit and maintenance # -------------------------------------------------------------------------- def get_audit_logs( self, api_key: str, start_date: Optional[datetime] = None, end_date: Optional[datetime] = None, limit: int = 100, ) -> List[Dict[str, Any]]: """Retrieve audit logs for a given API key.""" query = "SELECT * FROM usage_log WHERE api_key = ?" params = [api_key] if start_date: query += " AND timestamp >= ?" params.append(start_date.timestamp()) if end_date: query += " AND timestamp <= ?" params.append(end_date.timestamp()) query += " ORDER BY timestamp DESC LIMIT ?" params.append(limit) with self._get_conn() as conn: rows = conn.execute(query, params).fetchall() return [dict(row) for row in rows] def clean_old_logs(self): """Delete logs older than retention period for each tier, and old idempotency keys.""" with self._get_conn() as conn: # Delete old usage logs for tier in Tier: retention_days = tier.audit_log_retention_days if retention_days is None: continue cutoff = time.time() - retention_days * 86400 conn.execute( "DELETE FROM usage_log WHERE tier = ? AND timestamp < ?", (tier.value, cutoff) ) # Delete idempotency keys older than 7 days cutoff = time.time() - 7 * 86400 conn.execute( "DELETE FROM idempotency_keys WHERE consumed_at < ?", (cutoff,)) conn.commit() # -------------------------------------------------------------------------- # Global instance and FastAPI dependency (fail‑closed) # -------------------------------------------------------------------------- tracker: Optional[UsageTracker] = None def init_tracker( db_path: str = "arf_usage.db", redis_url: Optional[str] = None): """Initialize the global tracker. Must be called before enforce_quota.""" global tracker tracker = UsageTracker(db_path, redis_url) def update_key_tier(api_key: str, new_tier: Tier) -> bool: """Globally accessible helper to update API key tier.""" if tracker is None: return False return tracker.update_api_key_tier(api_key, new_tier) async def enforce_quota(request: Request, api_key: str = None): """ Dependency that checks API key and remaining quota. FAILS CLOSED: if tracker not initialised, raises HTTP 503. """ # P0 fix: No fallback that allows all requests if tracker is None: raise HTTPException( status_code=503, detail="Usage tracking service not initialised. Please contact administrator.") # Extract API key from header or query if api_key is None: auth_header = request.headers.get("Authorization") if auth_header and auth_header.startswith("Bearer "): api_key = auth_header[7:] else: api_key = request.query_params.get("api_key") if not api_key: raise HTTPException(status_code=401, detail="Missing API key") tier = tracker.get_tier(api_key) if tier is None: raise HTTPException( status_code=403, detail="Invalid or inactive API key") remaining = tracker.get_remaining_quota(api_key, tier) if remaining is not None and remaining <= 0: raise HTTPException(status_code=429, detail="Monthly evaluation quota exceeded") # Store in request state for later logging (optional) request.state.api_key = api_key request.state.tier = tier return {"api_key": api_key, "tier": tier, "remaining": remaining}