Spaces:

hdremover
/

hd-background-remover

Running

hd-background-remover / postgres_store.py

Update postgres_store.py

42a645d verified about 2 months ago

23.1 kB

"""
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"],
    }

	"""
	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"],
	}