Create postgres_store.py

postgres_store.py

@@ -0,0 +1,222 @@
+"""
+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 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}")