Upload folder using huggingface_hub

alembic/versions/a1f3c9d2e6b7_create_api_keys_table.py

@@ -0,0 +1,58 @@
+"""create api_keys table (moves API key storage off ephemeral SQLite)
+
+api_keys previously lived only in a SQLite file per service (arf_usage.db),
+which is wiped on every Render deploy/restart on the Free plan and was also
+independently duplicated between arf-api and arf-gateway. This creates the
+durable, single source of truth in Postgres. Rows are hashed at rest
+(pepper-HMAC lookup_hash + salted key_hash) -- no plaintext key column,
+since there is no pre-pepper data to migrate here (confirmed with the user:
+existing SQLite api_keys rows in both services are safe to discard, keys get
+reissued via POST /admin/keys).
+
+Revision ID: a1f3c9d2e6b7
+Revises: d36deffe7fa2
+Create Date: 2026-08-24 00:00:00.000000
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'a1f3c9d2e6b7'
+down_revision: Union[str, Sequence[str], None] = 'd36deffe7fa2'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # No FK to tenants.id here, deliberately -- the pre-existing SQLite
+    # schema this replaces never enforced one either (seeded/demo keys via
+    # ARF_API_KEYS may reference synthetic tenant_ids that don't have a
+    # tenants row), and adding one now would be a behavior change beyond
+    # this migration's scope.
+    op.create_table(
+        'api_keys',
+        sa.Column('id', sa.Integer(), nullable=False),
+        sa.Column('tenant_id', sa.String(length=64), nullable=False),
+        sa.Column('tier', sa.String(length=32), nullable=False),
+        sa.Column('created_at', sa.DateTime(), nullable=False),
+        sa.Column('last_used_at', sa.DateTime(), nullable=True),
+        sa.Column('is_active', sa.Boolean(), nullable=False, server_default=sa.true()),
+        sa.Column('salt', sa.String(length=64), nullable=False),
+        sa.Column('key_hash', sa.String(length=64), nullable=False),
+        sa.Column('lookup_hash', sa.String(length=64), nullable=False),
+        sa.PrimaryKeyConstraint('id'),
+    )
+    op.create_index(op.f('ix_api_keys_tenant_id'), 'api_keys', ['tenant_id'], unique=False)
+    op.create_unique_constraint('uq_api_keys_lookup_hash', 'api_keys', ['lookup_hash'])
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    op.drop_constraint('uq_api_keys_lookup_hash', 'api_keys', type_='unique')
+    op.drop_index(op.f('ix_api_keys_tenant_id'), table_name='api_keys')
+    op.drop_table('api_keys')

app/api/routes_admin.py

@@ -82,19 +82,20 @@ async def list_api_keys(limit: int = 100, offset: int = 0):
     longer has -- query `/admin/keys/{api_key}/audit` with the real key for
     per-key usage/audit history instead.
     """
-    with tracker._get_conn() as conn:
-        rows = conn.execute(
+    with tracker._get_pg_conn() as conn:
+        rows = tracker._pg_execute(
+            conn,
             "SELECT lookup_hash, tier, created_at, last_used_at, is_active FROM api_keys "
-            "WHERE lookup_hash != '' ORDER BY created_at DESC LIMIT ? OFFSET ?",
+            "ORDER BY created_at DESC LIMIT %s OFFSET %s",
             (limit, offset)
         ).fetchall()
+        conn.commit()
         keys = [
             {
                 "key_id": row["lookup_hash"],
                 "tier": row["tier"],
-                "created_at": datetime.fromtimestamp(row["created_at"]).isoformat(),
-                "last_used_at": datetime.fromtimestamp(row["last_used_at"]).isoformat()
-                if row["last_used_at"] else None,
+                "created_at": row["created_at"].isoformat(),
+                "last_used_at": row["last_used_at"].isoformat() if row["last_used_at"] else None,
                 "is_active": bool(row["is_active"]),
             }
             for row in rows
@@ -110,13 +111,14 @@ async def update_key_tier(
     if req.tier not in [t.value for t in Tier]:
         raise HTTPException(
             status_code=400, detail=f"Invalid tier. Must be one of {[t.value for t in Tier]}")
-    with tracker._get_conn() as conn:
-        row = conn.execute(
-            "SELECT lookup_hash FROM api_keys WHERE lookup_hash = ?", (key_id,)).fetchone()
+    with tracker._get_pg_conn() as conn:
+        row = tracker._pg_execute(
+            conn, "SELECT lookup_hash FROM api_keys WHERE lookup_hash = %s", (key_id,)).fetchone()
         if not row:
+            conn.rollback()
             raise HTTPException(status_code=404, detail="API key not found")
-        conn.execute("UPDATE api_keys SET tier = ? WHERE lookup_hash = ?",
-                     (req.tier, key_id))
+        tracker._pg_execute(
+            conn, "UPDATE api_keys SET tier = %s WHERE lookup_hash = %s", (req.tier, key_id))
         conn.commit()
     return {"message": f"Tier updated to {req.tier}"}
 
@@ -124,13 +126,14 @@ async def update_key_tier(
 @router.delete("/keys/{key_id}", dependencies=[Depends(verify_admin)])
 async def deactivate_api_key(
         key_id: str = Path(..., description="The key_id from GET /admin/keys (not the raw API key)")):
-    with tracker._get_conn() as conn:
-        row = conn.execute(
-            "SELECT lookup_hash FROM api_keys WHERE lookup_hash = ?", (key_id,)).fetchone()
+    with tracker._get_pg_conn() as conn:
+        row = tracker._pg_execute(
+            conn, "SELECT lookup_hash FROM api_keys WHERE lookup_hash = %s", (key_id,)).fetchone()
         if not row:
+            conn.rollback()
             raise HTTPException(status_code=404, detail="API key not found")
-        conn.execute(
-            "UPDATE api_keys SET is_active = 0 WHERE lookup_hash = ?", (key_id,))
+        tracker._pg_execute(
+            conn, "UPDATE api_keys SET is_active = false WHERE lookup_hash = %s", (key_id,))
         conn.commit()
     return {"message": "API key deactivated"}
 
@@ -150,9 +153,11 @@ async def get_audit_logs(
 
 @router.get("/stats", dependencies=[Depends(verify_admin)])
 async def get_global_stats():
+    with tracker._get_pg_conn() as pg_conn:
+        total_keys = tracker._pg_execute(
+            pg_conn, "SELECT COUNT(*) FROM api_keys WHERE is_active = true").fetchone()["count"]
+        pg_conn.commit()
     with tracker._get_conn() as conn:
-        total_keys = conn.execute(
-            "SELECT COUNT(*) FROM api_keys WHERE is_active = 1").fetchone()[0]
         total_requests = conn.execute(
             "SELECT COUNT(*) FROM usage_log").fetchone()[0]
         by_tier = conn.execute(

app/api/routes_payments.py

@@ -50,6 +50,13 @@ async def create_checkout_session(req: CheckoutRequest):
             cancel_url=req.cancel_url,
             metadata={"api_key": req.api_key},
             client_reference_id=req.api_key,
+            # checkout.session.completed carries this metadata via
+            # session.metadata (handled below), but customer.subscription.*
+            # events only carry the *subscription's own* metadata -- Stripe
+            # does not copy Session.metadata onto the Subscription it
+            # creates. Without this, cancellations can't be traced back to
+            # an api_key and PRO tier never downgrades.
+            subscription_data={"metadata": {"api_key": req.api_key}},
         )
         return {"sessionId": checkout_session.id, "url": checkout_session.url}
     except Exception as e:

app/core/usage_tracker.py

@@ -10,11 +10,15 @@ HMAC-SHA256(pepper, key) -- a deterministic but one-way value computed with
 a server-only secret (ARF_KEY_PEPPER), so a leaked database alone does not
 expose usable keys. A per-row random salt plus a second SHA-256 check is a
 defense-in-depth verification layer after the row is found by lookup hash.
-This mirrors arf-gateway's internal/auth/apikey.go, which both services may
-point at the same underlying database file for -- same column names and
-hash algorithm, so rows created by either service validate on both. Only
-this table is affected: `monthly_counts` and `usage_log` are unchanged and
-still keyed by the raw API key, since arf-gateway reads them that way.
+This mirrors arf-gateway's internal/auth/apikey.go.
+
+`api_keys` itself lives in Postgres (DATABASE_URL), not in the SQLite file
+the rest of this module uses -- a single durable table that both arf-api
+and arf-gateway point at, instead of each service's own local SQLite copy
+(which on Render's Free plan is wiped on every deploy/restart anyway).
+`monthly_counts`, `usage_log`, and `idempotency_keys` are unchanged and
+still SQLite, still keyed by the raw API key, since arf-gateway reads
+`monthly_counts` that way.
 """
 import hashlib
 import hmac
@@ -24,6 +28,9 @@ import secrets
 import sqlite3
 import threading
 import time
+
+import psycopg2
+import psycopg2.extras
 from contextlib import contextmanager
 from datetime import datetime, timedelta
 from dataclasses import dataclass
@@ -104,9 +111,14 @@ class UsageTracker:
                 f"ARF_KEY_PEPPER is too short ({len(self._pepper)} chars); "
                 "use at least 32 random characters."
             )
+        self._pg_dsn = os.getenv("DATABASE_URL", "")
+        if not self._pg_dsn:
+            raise RuntimeError(
+                "DATABASE_URL is not set -- refusing to start without it, "
+                "since api_keys is stored in Postgres, not SQLite."
+            )
         self._local = threading.local()
         self._init_db()
-        self._migrate_plaintext_keys()
 
         self._redis_client = None
         if redis_url and REDIS_AVAILABLE:
@@ -125,13 +137,15 @@ class UsageTracker:
         already been found via lookup hash."""
         return hashlib.sha256(bytes.fromhex(salt_hex) + key.encode()).hexdigest()
 
-    def _verify_key(self, conn, api_key: str) -> Optional[sqlite3.Row]:
+    def _verify_key(self, conn, api_key: str) -> Optional[dict]:
         """Look up a row by pepper-HMAC, then verify with the salted hash.
-        Returns the row (tenant_id, tier, is_active, salt, key_hash) if the
-        key is valid and active, else None."""
-        row = conn.execute(
+        `conn` is a Postgres connection from _get_pg_conn. Returns the row
+        (tenant_id, tier, is_active, salt, key_hash) if the key is valid and
+        active, else None."""
+        row = self._pg_execute(
+            conn,
             "SELECT tenant_id, tier, is_active, salt, key_hash FROM api_keys "
-            "WHERE lookup_hash = ?",
+            "WHERE lookup_hash = %s",
             (self._lookup_hash(api_key),)
         ).fetchone()
         if not row or not row["is_active"]:
@@ -142,7 +156,10 @@ class UsageTracker:
 
     @contextmanager
     def _get_conn(self):
-        """Get a thread‑local SQLite connection with WAL and immediate transactions."""
+        """Get a thread‑local SQLite connection with WAL and immediate transactions.
+
+        Backs usage_log/monthly_counts/idempotency_keys only -- api_keys
+        lives in Postgres, see _get_pg_conn below."""
         if not hasattr(self._local, "conn"):
             self._local.conn = sqlite3.connect(
                 self.db_path, check_same_thread=False, isolation_level=None)
@@ -150,38 +167,65 @@ class UsageTracker:
             self._local.conn.execute("PRAGMA journal_mode=WAL")
         yield self._local.conn
 
-    def _init_db(self):
-        """Initialise SQLite tables with tenant_id support."""
-        with self._get_conn() as conn:
-            # `key` is kept only for backward compatibility with pre-pepper
-            # databases during migration below -- new rows never populate
-            # it with a real value (see get_or_create_api_key).
-            conn.execute("""
+    @contextmanager
+    def _get_pg_conn(self):
+        """Get a thread-local Postgres connection for the api_keys table.
+        Rows come back as dict-like objects (row["col"]) via RealDictCursor,
+        matching the sqlite3.Row access pattern used elsewhere in this file."""
+        if not hasattr(self._local, "pg_conn") or self._local.pg_conn.closed:
+            self._local.pg_conn = psycopg2.connect(
+                self._pg_dsn, cursor_factory=psycopg2.extras.RealDictCursor)
+        yield self._local.pg_conn
+
+    @staticmethod
+    def _pg_execute(conn, sql: str, params: tuple = ()):
+        """Run a query against a Postgres connection and return the cursor,
+        so callers can chain .fetchone()/.fetchall() the same way sqlite3's
+        conn.execute(...) is used elsewhere in this file."""
+        cur = conn.cursor()
+        cur.execute(sql, params)
+        return cur
+
+    def _init_pg_db(self):
+        """Idempotently ensure the Postgres api_keys table/index exist.
+
+        The canonical schema is the Alembic migration
+        (alembic/versions/*_create_api_keys_table.py) -- but nothing in this
+        codebase runs `alembic upgrade head` automatically on deploy (a
+        known gap, tracked separately), and arf-gateway's Go code needs the
+        same table without going through Python/Alembic at all. Mirroring
+        the same CREATE TABLE IF NOT EXISTS self-healing pattern this file
+        already uses for its SQLite tables keeps both services (and tests)
+        working whether or not the migration has actually been applied.
+        Column set/types must stay in sync with that migration."""
+        with self._get_pg_conn() as conn:
+            self._pg_execute(conn, """
                 CREATE TABLE IF NOT EXISTS api_keys (
-                    key TEXT PRIMARY KEY,
-                    tenant_id TEXT NOT NULL,
-                    tier TEXT NOT NULL,
-                    created_at REAL NOT NULL,
-                    last_used_at REAL,
-                    is_active INTEGER DEFAULT 1,
-                    salt TEXT NOT NULL DEFAULT '',
-                    key_hash TEXT NOT NULL DEFAULT '',
-                    lookup_hash TEXT NOT NULL DEFAULT ''
+                    id SERIAL PRIMARY KEY,
+                    tenant_id VARCHAR(64) NOT NULL,
+                    tier VARCHAR(32) NOT NULL,
+                    created_at TIMESTAMP NOT NULL,
+                    last_used_at TIMESTAMP,
+                    is_active BOOLEAN NOT NULL DEFAULT true,
+                    salt VARCHAR(64) NOT NULL,
+                    key_hash VARCHAR(64) NOT NULL,
+                    lookup_hash VARCHAR(64) NOT NULL
                 )
             """)
-            # Idempotent: add security columns to a pre-existing table that
-            # predates this fix.
-            existing_cols = {
-                row["name"] for row in conn.execute("PRAGMA table_info(api_keys)").fetchall()
-            }
-            for col in ("salt", "key_hash", "lookup_hash"):
-                if col not in existing_cols:
-                    conn.execute(f"ALTER TABLE api_keys ADD COLUMN {col} TEXT NOT NULL DEFAULT ''")
-            conn.execute(
-                "CREATE UNIQUE INDEX IF NOT EXISTS idx_api_keys_lookup_hash "
-                "ON api_keys(lookup_hash) WHERE lookup_hash != ''"
-            )
+            self._pg_execute(conn, """
+                CREATE UNIQUE INDEX IF NOT EXISTS uq_api_keys_lookup_hash
+                ON api_keys (lookup_hash)
+            """)
+            self._pg_execute(conn, """
+                CREATE INDEX IF NOT EXISTS ix_api_keys_tenant_id
+                ON api_keys (tenant_id)
+            """)
             conn.commit()
+
+    def _init_db(self):
+        """Initialise SQLite tables for usage_log/monthly_counts/idempotency_keys."""
+        self._init_pg_db()
+        with self._get_conn() as conn:
             conn.execute("""
                 CREATE TABLE IF NOT EXISTS usage_log (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -219,35 +263,6 @@ class UsageTracker:
     def _get_month_key(self) -> str:
         return datetime.now().strftime("%Y-%m")
 
-    def _migrate_plaintext_keys(self):
-        """Find rows left over from before pepper-HMAC lookup existed
-        (lookup_hash empty, key non-empty), hash them, and clear the
-        plaintext key column. Safe to run on every startup: already-migrated
-        rows (key already cleared) are no-ops."""
-        with self._get_conn() as conn:
-            legacy = conn.execute(
-                "SELECT key FROM api_keys WHERE lookup_hash = '' AND key != ''"
-            ).fetchall()
-            if not legacy:
-                return
-            import logging
-            logger = logging.getLogger(__name__)
-            for row in legacy:
-                plaintext_key = row["key"]
-                salt = secrets.token_hex(16)
-                conn.execute(
-                    "UPDATE api_keys SET lookup_hash = ?, salt = ?, key_hash = ?, key = NULL "
-                    "WHERE key = ?",
-                    (self._lookup_hash(plaintext_key), salt,
-                     self._salted_hash(plaintext_key, salt), plaintext_key)
-                )
-            conn.commit()
-            logger.warning(
-                "Migrated %d plaintext API key(s) to pepper-HMAC lookup; this database "
-                "previously stored keys in plaintext -- rotate keys and check backups/"
-                "exports for leaked copies.", len(legacy)
-            )
-
     def get_or_create_api_key(self, key: str, tenant_id: str, tier: Tier = Tier.FREE) -> bool:
         """
         Register a new API key for a given tenant.
@@ -261,25 +276,24 @@ class UsageTracker:
             True if key was created (or already exists for the same tenant).
         """
         lookup_hash = self._lookup_hash(key)
-        with self._get_conn() as conn:
-            row = conn.execute(
-                "SELECT tenant_id FROM api_keys WHERE lookup_hash = ?", (lookup_hash,)
+        with self._get_pg_conn() as conn:
+            row = self._pg_execute(
+                conn, "SELECT tenant_id FROM api_keys WHERE lookup_hash = %s", (lookup_hash,)
             ).fetchone()
             if row:
                 # Key already exists – ensure it belongs to the same tenant
                 if row["tenant_id"] != tenant_id:
+                    conn.rollback()
                     raise ValueError(f"Key {key[:8]}... already belongs to a different tenant.")
+                conn.commit()
                 return True
             salt = secrets.token_hex(16)
-            conn.execute(
-                # `key` is left NULL (not '') so a second row can be created
-                # without violating the PRIMARY KEY constraint on that
-                # column -- SQLite permits multiple NULLs in a non-INTEGER
-                # PRIMARY KEY column, but not multiple empty strings.
+            self._pg_execute(
+                conn,
                 "INSERT INTO api_keys "
-                "(key, tenant_id, tier, created_at, is_active, salt, key_hash, lookup_hash) "
-                "VALUES (NULL, ?, ?, ?, ?, ?, ?, ?)",
-                (tenant_id, tier.value, time.time(), 1,
+                "(tenant_id, tier, created_at, is_active, salt, key_hash, lookup_hash) "
+                "VALUES (%s, %s, %s, %s, %s, %s, %s)",
+                (tenant_id, tier.value, datetime.utcnow(), True,
                  salt, self._salted_hash(key, salt), lookup_hash)
             )
             conn.commit()
@@ -287,27 +301,28 @@ class UsageTracker:
 
     def get_tier(self, api_key: str) -> Optional[Tier]:
         """Return the tier for a given API key, or None if key invalid/inactive."""
-        with self._get_conn() as conn:
+        with self._get_pg_conn() as conn:
             row = self._verify_key(conn, api_key)
             return Tier(row["tier"]) if row else None
 
     def get_tenant_id(self, api_key: str) -> Optional[str]:
         """Return the tenant ID associated with the API key, or None if key invalid."""
-        with self._get_conn() as conn:
+        with self._get_pg_conn() as conn:
             row = self._verify_key(conn, api_key)
             return row["tenant_id"] if row else None
 
     def update_api_key_tier(self, api_key: str, new_tier: Tier) -> bool:
         """Update the tier of an existing API key. Returns True if successful."""
         lookup_hash = self._lookup_hash(api_key)
-        with self._get_conn() as conn:
-            row = conn.execute(
-                "SELECT lookup_hash FROM api_keys WHERE lookup_hash = ?", (lookup_hash,)
+        with self._get_pg_conn() as conn:
+            row = self._pg_execute(
+                conn, "SELECT lookup_hash FROM api_keys WHERE lookup_hash = %s", (lookup_hash,)
             ).fetchone()
             if not row:
+                conn.rollback()
                 return False
-            conn.execute(
-                "UPDATE api_keys SET tier = ? WHERE lookup_hash = ?",
+            self._pg_execute(
+                conn, "UPDATE api_keys SET tier = %s WHERE lookup_hash = %s",
                 (new_tier.value, lookup_hash))
             conn.commit()
             return True

tests/test_routes_admin.py

@@ -0,0 +1,83 @@
+"""
+Route-level tests for the /admin/keys endpoints against the real,
+Postgres-backed UsageTracker -- these routes are normally exercised through
+the module-level `tracker` singleton (conftest.py replaces it globally with
+MockTracker for every other test), so both `tracker` and `ADMIN_API_KEY` are
+monkeypatched directly on `app.api.routes_admin` for the duration of this
+module. This is the main regression coverage for the api_keys-on-Postgres
+migration: it exercises the exact raw SQL routes_admin.py runs against
+`tracker._get_pg_conn()`.
+"""
+import hashlib
+import hmac
+import os
+
+import pytest
+
+from app.core.usage_tracker import UsageTracker
+import app.api.routes_admin as routes_admin
+
+TEST_ADMIN_KEY = "test-admin-key-for-routes-admin-tests"
+TEST_PEPPER = os.environ["ARF_KEY_PEPPER"]  # set in conftest.py before app import
+
+
+def _key_id(raw_key: str) -> str:
+    """Reproduce UsageTracker._lookup_hash without a tracker instance, so
+    tests can locate the row created for a given raw key deterministically."""
+    return hmac.new(TEST_PEPPER.encode(), raw_key.encode(), hashlib.sha256).hexdigest()
+
+
+@pytest.fixture(autouse=True)
+def real_tracker(monkeypatch):
+    real = UsageTracker(db_path=":memory:")
+    monkeypatch.setattr(routes_admin, "tracker", real)
+    monkeypatch.setattr(routes_admin, "ADMIN_API_KEY", TEST_ADMIN_KEY)
+    yield real
+
+
+def test_create_list_update_deactivate_key(client):
+    create_resp = client.post(
+        "/admin/keys",
+        params={"admin_key": TEST_ADMIN_KEY},
+        json={"tier": "free", "org_name": "Test Org"},
+    )
+    assert create_resp.status_code == 200
+    body = create_resp.json()
+    api_key = body["api_key"]
+    assert body["tier"] == "free"
+    key_id = _key_id(api_key)
+
+    list_resp = client.get("/admin/keys", params={"admin_key": TEST_ADMIN_KEY})
+    assert list_resp.status_code == 200
+    keys_by_id = {row["key_id"]: row for row in list_resp.json()["keys"]}
+    assert key_id in keys_by_id
+    assert keys_by_id[key_id]["tier"] == "free"
+    assert keys_by_id[key_id]["is_active"] is True
+
+    patch_resp = client.patch(
+        f"/admin/keys/{key_id}/tier",
+        params={"admin_key": TEST_ADMIN_KEY},
+        json={"tier": "pro"},
+    )
+    assert patch_resp.status_code == 200
+
+    list_resp2 = client.get("/admin/keys", params={"admin_key": TEST_ADMIN_KEY})
+    assert list_resp2.json()["keys"][0]  # non-empty, sanity check
+    keys_by_id2 = {row["key_id"]: row for row in list_resp2.json()["keys"]}
+    assert keys_by_id2[key_id]["tier"] == "pro"
+
+    delete_resp = client.delete(f"/admin/keys/{key_id}", params={"admin_key": TEST_ADMIN_KEY})
+    assert delete_resp.status_code == 200
+
+    list_resp3 = client.get("/admin/keys", params={"admin_key": TEST_ADMIN_KEY})
+    keys_by_id3 = {row["key_id"]: row for row in list_resp3.json()["keys"]}
+    assert keys_by_id3[key_id]["is_active"] is False
+
+
+def test_update_nonexistent_key_returns_404(client):
+    resp = client.patch(
+        "/admin/keys/does-not-exist/tier",
+        params={"admin_key": TEST_ADMIN_KEY},
+        json={"tier": "pro"},
+    )
+    assert resp.status_code == 404