""" postgres_store.py Postgres-backed credits/API-key layer for HD Remover's HuggingFace backend. Replaces the JSON-based credit/usage tracking in data_store.py (calls_today, daily limit, users.json) with queries against the SAME Postgres database the HDRemover website uses — see migrations/002_add_website_login_columns.sql and migrations/003_add_dashboard_columns.sql in the website repo for the exact schema this depends on. WHY A SEPARATE FILE RATHER THAN EDITING data_store.py IN PLACE: data_store.py's JSON system still owns things this file does NOT touch — plans.json (admin-editable plan limits), the Whop-to-plan mapping, and the webhook audit log. Those are being removed/replaced separately (Whop is being dropped entirely in favor of Safepay, which bills through the website, not this backend). Keeping this as a new file makes that later cleanup a deletion of whole functions/files rather than a tangle of half-migrated code in one already-large file. WHAT THIS FILE OWNS: - API key verification (bcrypt hash comparison against api_keys table) - Credit balance checks and deduction (monthly_credits first, then lifetime_credits — matches the website dashboard's display order and the agreed business rule: subscription credits are "use it or lose it" each cycle, so they should be spent first) - last_used_at tracking on the api_keys row (surfaced in the website dashboard's API Keys tab) WHAT THIS FILE DELIBERATELY DOES NOT DO: - No plan/daily-limit concept. There is no "free/starter/pro/master" plan tier here — a request is authorized purely by "does this API key belong to a user with credits > 0". Per-minute rate limiting (abuse prevention, not billing) still lives in app.py's in-memory _minute_windows — that's a different concern from credits and doesn't need to move to Postgres. - No user creation. Accounts are only ever created by the website (email/password, Google, or Facebook signup) — this backend only ever reads/updates an existing user's credit balance, never inserts a new user row. Env var required: DATABASE_URL — same Postgres connection string used by the website (Supabase Session Pooler connection string, NOT the direct connection — see the website's lib/db.ts for why: this backend, like the website, holds a long-lived connection pool rather than one-shot connections per request). """ import os import logging import threading from datetime import datetime, timezone import bcrypt import psycopg2 import psycopg2.pool from psycopg2.extras import RealDictCursor logger = logging.getLogger("hd_remover.postgres_store") DATABASE_URL = os.environ.get("DATABASE_URL", "") # A small connection pool (not a single global connection) so concurrent # requests don't serialize on one connection — mirrors the website's own # lib/db.ts, which uses a pg.Pool with max: 10 for the same reason. _pool = None _pool_lock = threading.Lock() def _get_pool(): global _pool if _pool is not None: return _pool with _pool_lock: if _pool is None: if not DATABASE_URL: raise RuntimeError( "DATABASE_URL environment variable is not set. This backend " "cannot verify API keys or check credits without it." ) _pool = psycopg2.pool.ThreadedConnectionPool(minconn=1, maxconn=10, dsn=DATABASE_URL) logger.info("Postgres connection pool created") return _pool class _PooledConnection: """Context manager that borrows a connection from the pool and always returns it, even if the query inside raises.""" def __enter__(self): self._pool = _get_pool() self._conn = self._pool.getconn() return self._conn def __exit__(self, exc_type, exc_val, exc_tb): if exc_type is not None: # Roll back so a failed query doesn't leave the connection # mid-transaction when it's returned to the pool for reuse. self._conn.rollback() self._pool.putconn(self._conn) def utcnow(): return datetime.now(timezone.utc) # ── API key verification ──────────────────────────────────────────────────── def find_user_by_api_key(raw_api_key: str): """ Verifies a raw API key against the api_keys table's bcrypt hashes and returns the owning user's credit info, or None if the key is invalid. IMPORTANT — this is O(number of API keys in the system), not O(1): bcrypt hashes can't be looked up by an indexed equality match (the whole point of bcrypt is that the same input never produces the same hash twice, via its per-hash salt), so there is no "SELECT ... WHERE key_hash = ?" shortcut. Every active key's hash must be checked with bcrypt.checkpw() until one matches. This uses the key_prefix column (see migrations/003) to narrow the candidate set first — key_prefix is derived deterministically from the raw key (not salted), so it CAN be indexed and equality-matched, cutting the bcrypt comparison down to (usually) a single row instead of scanning every key in the table. """ if not raw_api_key or not raw_api_key.startswith("hdrm_"): return None # Matches the website's generateRawKey() in # app/api/dashboard/api-keys/route.ts: prefix is "hdrm_" + first 8 hex # chars of the random part, 13 chars total. key_prefix = raw_api_key[:13] with _PooledConnection() as conn: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( """ SELECT ak.id AS api_key_id, ak.key_hash, ak.uid, u.plan_id, u.monthly_credits, u.lifetime_credits FROM api_keys ak JOIN users u ON u.uid = ak.uid WHERE ak.key_prefix = %s """, (key_prefix,), ) candidates = cur.fetchall() for row in candidates: stored_hash = row["key_hash"] if bcrypt.checkpw(raw_api_key.encode("utf-8"), stored_hash.encode("utf-8")): return { "api_key_id": row["api_key_id"], "uid": row["uid"], "plan_id": row["plan_id"], "monthly_credits": row["monthly_credits"], "lifetime_credits": row["lifetime_credits"], } return None # ── Credit checking + deduction ───────────────────────────────────────────── def has_credits(user_info: dict) -> bool: """Returns True if the user has at least 1 credit available, from either bucket combined.""" return (user_info.get("monthly_credits", 0) + user_info.get("lifetime_credits", 0)) > 0 def deduct_one_credit(uid: str) -> bool: """ Atomically deducts exactly 1 credit from a user's balance — monthly_credits first, falling back to lifetime_credits only once monthly is at 0 (see the module docstring for why this order was chosen). Uses a single UPDATE with a CASE expression rather than a read-then-write (SELECT balance, check in Python, UPDATE) — this avoids a race condition where two concurrent requests both read "1 credit left" before either writes, and both proceed, resulting in -1 credits. The database performs the check-and-decrement as one atomic operation instead. Returns True if a credit was actually deducted, False if the user had zero credits in both buckets (nothing was deducted — the row is matched by the WHERE clause only when there's something to spend). """ with _PooledConnection() as conn: with conn.cursor() as cur: cur.execute( """ UPDATE users SET monthly_credits = CASE WHEN monthly_credits > 0 THEN monthly_credits - 1 ELSE monthly_credits END, lifetime_credits = CASE WHEN monthly_credits > 0 THEN lifetime_credits ELSE lifetime_credits - 1 END WHERE uid = %s AND (monthly_credits > 0 OR lifetime_credits > 0) """, (uid,), ) deducted = cur.rowcount > 0 conn.commit() return deducted def log_request(uid: str, api_key_id: str, feature: str, model_key: str, plan_id: str, credits_charged: int, gpu: str = None): """ Records one row in request_logs for the admin dashboard's analytics (per-model breakdown, daily counts, free-vs-paid volume). Added 2026-08-23 — see migrations/004_add_request_logs.sql for the table. Fire-and-forget in spirit, same as touch_api_key_last_used: this must NEVER block or fail an actual image-processing request that already succeeded. Any exception here is logged, not raised. """ try: with _PooledConnection() as conn: with conn.cursor() as cur: cur.execute( """ INSERT INTO request_logs (uid, api_key_id, feature, model_key, plan_id_at_request, is_free, credits_charged, gpu) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) """, (uid, api_key_id, feature, model_key, plan_id, plan_id == "free", credits_charged, gpu), ) conn.commit() except Exception as e: logger.warning(f"log_request failed for uid={uid} feature={feature}: {e}") def touch_api_key_last_used(api_key_id: str): """ Updates last_used_at on the api_keys row — surfaced in the website dashboard's API Keys tab so a person can tell which of their keys is actually in use (e.g. to safely identify an old/unused key before revoking it). Fire-and-forget in spirit (failure here should never block an actual image-processing request that already succeeded) — called after the real work is done, and any exception is logged, not raised. """ try: with _PooledConnection() as conn: with conn.cursor() as cur: cur.execute( "UPDATE api_keys SET last_used_at = %s WHERE id = %s", (utcnow(), api_key_id), ) conn.commit() except Exception as e: logger.warning(f"touch_api_key_last_used failed for {api_key_id}: {e}")