postgres_store.py

"""
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}")