""" 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) and the webhook audit log. Those are being removed/replaced separately. 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, for BOTH key types the system now supports: 1. Sub-keys — user-visible, dashboard-generated (api_keys table, bcrypt hash comparison). This now covers BOTH paid and free registered users — a free user's key is a row in the exact same api_keys table as a paid user's, distinguished only by that uid's users.plan_id (== "free"). No separate table/flag was introduced for this (2026-09-02 decision) — see has_credits()/deduct_one_credit() below for how the extra api_credits pool this enables is actually spent. 2. Master keys — NEVER shown to the user, generated once server-side on a user's first-ever paid purchase (see the website's app/api/webhooks/whop/route.ts). Stored directly on the users row (master_key_hash / master_key_prefix columns) — not in api_keys — since a user has exactly one, unlike sub-keys which can be regenerated/rotated freely. Free users never have a master key — only ever a sub-key. Per the confirmed design (2026-08-26): a sub-key carries NO credits of its own — it exists purely to identify which user is calling. Whether a request arrives via a sub-key or directly via a master key, the credits/plan checked are always the SAME users row (that's why find_user_by_api_key() already joins to users for plan_id/ monthly_credits/lifetime_credits/api_credits regardless of which key type matched — no separate "look up the master key's balance" step is needed, since sub-keys and their owning master key both resolve to the same uid's row). - Credit balance checks and deduction, across THREE buckets on the users row (2026-09-02 — added api_credits, see below): 1. monthly_credits — paid subscription allowance, resets each billing cycle 2. lifetime_credits — paid one-time/rollover credits 3. api_credits — free registered users' API allowance (50/ month, confirmed 2026-09-02). Deliberately a SEPARATE pool from a free user's normal dashboard free_credits (the 100/month used for the website's own upload flow) — a free user generating an API key does not eat into their dashboard allowance or vice versa. Initial grant of 50 happens on the website (app/api/dashboard/api-keys/route.ts, on a free user's first-ever key). The MONTHLY RESET, unlike paid monthly_credits, has no billing webhook to hang off of (a free user has no subscription event) — so this backend does it lazily, inline, in find_user_by_api_key() below, see _maybe_reset_api_credits() (2026-09-03). There is no cron job or scheduled task anywhere in this system for this — deliberate, since the user is hosted on Hostinger with no convenient cron infra, and a lazy per-request check costs nothing extra on the (overwhelming majority of) requests where no reset is due. Deduction order confirmed 2026-09-02: monthly_credits first, then lifetime_credits, and ONLY THEN api_credits — this applies even to a paid user who still happens to have api_credits left over (edge case: a user was free, accumulated api_credits, then upgraded to paid — their paid buckets are still spent first, matching the existing "paid credits first" business rule rather than carving out a special case for how api_credits got there). - last_used_at tracking on the api_keys row (surfaced in the website dashboard's API Keys tab) — master-key authentication does not touch this, since there is no api_keys row for a master key to update. WHAT THIS FILE DELIBERATELY DOES NOT DO: - No plan/daily-limit concept beyond the three credit buckets above. 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, and no api_keys row creation. Accounts are only ever created by the website (email/password, Google, or Facebook signup), and API keys (both free and paid) are only ever generated by the website's dashboard — this backend only ever reads/updates an existing user's credit balance and an existing key's last_used_at, never inserts a new user or key row. Master keys are likewise only ever generated by the website's webhook, never by this backend. Resetting api_credits to 50 each month is also the website's job (cron), not this file's. 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) # Free registered users' monthly API-credit allowance — separate pool # from their normal dashboard/upload free_credits, confirmed 2026-09-02. FREE_MONTHLY_API_CREDITS = 50 def _current_month_start(): now = utcnow() return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) def _maybe_reset_api_credits(uid: str, api_credits_reset_at): """ Lazy monthly reset for a free user's api_credits pool (2026-09-03 — see postgres_store.py's module docstring). There is no cron/webhook that renews api_credits the way Whop's payment.succeeded webhook renews monthly_credits for paid plans, since a free user never has a billing event to hang a renewal off of. Instead, this runs inline on every sub-key lookup and only issues an UPDATE when a reset is actually due (api_credits_reset_at is NULL, meaning it was never granted/reset yet, or falls before the start of the current calendar month) — for every other request this is a single in-memory datetime comparison against the row already fetched by the caller's query, no extra round trip. Only ever called for a SUB-KEY match in find_user_by_api_key() — not for master keys, since master keys are paid-only and never touch api_credits at all. Resets to 50 (FREE_MONTHLY_API_CREDITS below) rather than adding — matches the "replace, don't carry forward" rule already used for paid monthly_credits (see the website's Whop webhook), so unused api_credits from last month don't stack. Returns the credit value the caller should use for this request (the freshly-reset 50 if a reset just happened, otherwise the value passed in unchanged) — callers must use this return value instead of their own stale api_credits read when one is available. """ now = utcnow() due = api_credits_reset_at is None or api_credits_reset_at < _current_month_start() if not due: return None # no reset happened — caller keeps its existing value with _PooledConnection() as conn: with conn.cursor() as cur: cur.execute( "UPDATE users SET api_credits = %s, api_credits_reset_at = %s WHERE uid = %s", (FREE_MONTHLY_API_CREDITS, now, uid), ) conn.commit() return FREE_MONTHLY_API_CREDITS # ── API key verification ──────────────────────────────────────────────────── def find_user_by_api_key(raw_api_key: str): """ Verifies a raw API key against EITHER the api_keys table (sub-keys — now used by BOTH free and paid registered users, see module docstring) OR the users table's master_key_hash column (master keys, paid-only), and returns the owning user's credit info — or None if the key matches neither. Both key types share the same "hdrm_" + 64-hex-char format and the same 13-char prefix-first-then-bcrypt lookup strategy (see the docstring note on why prefix narrowing is needed — bcrypt hashes can't be indexed/equality-matched directly). Checking sub-keys first is an arbitrary ordering choice, not a security requirement: a given raw key can only ever match one of the two tables, since each is generated independently and collision is cryptographically negligible. Callers (app.py) do not need to know or care which table actually matched, or whether the owning user is free or paid — the returned dict shape is identical in every case, and now always includes api_credits alongside monthly_credits/lifetime_credits so callers can check/spend all three buckets uniformly. The only difference "under the hood" is that a sub-key match also carries an api_key_id (used elsewhere for last_used_at tracking), while a master-key match does not, since there is no api_keys row involved. """ 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. Master keys use the exact # same format (see generateMasterKey() in the website's # app/api/webhooks/whop/route.ts), so the same prefix slice applies. key_prefix = raw_api_key[:13] # ── Try sub-keys first (api_keys table — free or paid owner) ───── 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, u.api_credits, u.api_credits_reset_at FROM api_keys ak JOIN users u ON u.uid = ak.uid WHERE ak.key_prefix = %s """, (key_prefix,), ) sub_key_candidates = cur.fetchall() for row in sub_key_candidates: stored_hash = row["key_hash"] if bcrypt.checkpw(raw_api_key.encode("utf-8"), stored_hash.encode("utf-8")): api_credits = row["api_credits"] # Lazy monthly reset — only free-plan rows ever accumulate # api_credits_reset_at drift worth checking; skipped for paid # users as a cheap no-op guard (a paid user's api_credits is # always 0 and irrelevant to routing/deduction anyway). if row["plan_id"] == "free": reset_value = _maybe_reset_api_credits(row["uid"], row["api_credits_reset_at"]) if reset_value is not None: api_credits = reset_value 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"], "api_credits": api_credits, } # ── Not a sub-key — try master keys (users.master_key_hash) ────── # Paid-only path: free users never have a master key, so api_credits # is still selected here for shape-consistency but will only ever be # relevant to a sub-key match in practice. with _PooledConnection() as conn: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( """ SELECT uid, master_key_hash, plan_id, monthly_credits, lifetime_credits, api_credits FROM users WHERE master_key_prefix = %s """, (key_prefix,), ) master_key_candidates = cur.fetchall() for row in master_key_candidates: stored_hash = row["master_key_hash"] if stored_hash and bcrypt.checkpw(raw_api_key.encode("utf-8"), stored_hash.encode("utf-8")): return { "api_key_id": None, # no api_keys row exists for a master-key match "uid": row["uid"], "plan_id": row["plan_id"], "monthly_credits": row["monthly_credits"], "lifetime_credits": row["lifetime_credits"], "api_credits": row["api_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 ANY of the three buckets combined: monthly_credits, lifetime_credits (both paid), or api_credits (free registered users' separate 50/ month API allowance, added 2026-09-02 — see module docstring).""" return ( user_info.get("monthly_credits", 0) + user_info.get("lifetime_credits", 0) + user_info.get("api_credits", 0) ) > 0 def has_paid_credits(user_info: dict) -> bool: """ Returns True only if the user has paid credits (monthly or lifetime) — deliberately EXCLUDES api_credits. Used by app.py to decide model/GPU routing: paid credits -> best model on Modal GPU, otherwise (including a free user spending api_credits) -> quality model on local HF CPU. Kept as its own function rather than inlined in app.py so the "what counts as paid" definition lives in one place alongside has_credits() and deduct_one_credit(), which must all agree on it. """ 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, trying each bucket in order: monthly_credits, then lifetime_credits, then api_credits (2026-09-02 addition — see module docstring for the confirmed priority: paid buckets are always spent first, even for a user who has api_credits left over, so a free-turned-paid user's leftover api_credits are only touched once both paid buckets are at zero). Uses a single UPDATE with CASE expressions 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 a negative balance. The database performs the check-and-decrement as one atomic operation instead. Returns True if a credit was actually deducted from any bucket, False if the user had zero credits across all three (nothing was deducted — the row is matched by the WHERE clause only when there's something to spend). Note: this operates on `uid` regardless of whether the caller authenticated via a sub-key or a master key, and regardless of which bucket ends up being spent — both key types resolve to the same users row. """ 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 WHEN lifetime_credits > 0 THEN lifetime_credits - 1 ELSE lifetime_credits END, api_credits = CASE WHEN monthly_credits > 0 THEN api_credits WHEN lifetime_credits > 0 THEN api_credits ELSE api_credits - 1 END WHERE uid = %s AND (monthly_credits > 0 OR lifetime_credits > 0 OR api_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). api_key_id may be None here (a master-key-authenticated request has no api_keys row) — callers should pass whatever find_user_by_api_key() returned for that field, including None, without special-casing it. 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). Applies equally to free and paid users' sub-keys. Callers MUST skip calling this when api_key_id is None (a master-key match) — there is no api_keys row to update in that case. This function itself does not guard against a None id being passed (an UPDATE ... WHERE id = NULL simply matches zero rows and is a silent no-op either way), but callers should still check explicitly rather than relying on that, since generating a spurious query on every master-key request is needless overhead. 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}") def get_user_by_uid(uid: str): """ Direct uid lookup — no API-key/bcrypt involved. Used ONLY by the trusted internal-service auth path in app.py (a website server-to- server call authenticated by a shared internal secret, not by a per-user API key). Returns the same dict shape as find_user_by_api_key() so callers can treat both paths identically, with api_key_id always None. Returns None if the uid doesn't exist. """ with _PooledConnection() as conn: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( "SELECT uid, plan_id, monthly_credits, lifetime_credits, api_credits FROM users WHERE uid = %s", (uid,), ) row = cur.fetchone() if not row: return None return { "api_key_id": None, "uid": row["uid"], "plan_id": row["plan_id"], "monthly_credits": row["monthly_credits"], "lifetime_credits": row["lifetime_credits"], "api_credits": row["api_credits"], }