""" OmniParse AI - Database Layer Dual-backend persistence: - **Supabase** (PostgreSQL) when SUPABASE_URL and SUPABASE_KEY are set. - **SQLite** as the fallback (file-based, suitable for HF Spaces). All public functions return plain dicts / lists — never raw Row / tuple objects. """ import json import os import sqlite3 import time import uuid from config import USE_SUPABASE, supabase_client, SESSION_TTL_SECONDS # --------------------------------------------------------------------------- # SQLite setup # --------------------------------------------------------------------------- _DB_PATH = os.environ.get("SQLITE_DB_PATH", "omniparse.db") def _sqlite_conn() -> sqlite3.Connection: """Return a new SQLite connection with row_factory enabled.""" conn = sqlite3.connect(_DB_PATH, check_same_thread=False) conn.row_factory = sqlite3.Row conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA foreign_keys=ON") return conn def _row_to_dict(row) -> dict: """Convert a sqlite3.Row (or similar) to a plain dict.""" if row is None: return None if isinstance(row, dict): return row return dict(row) # --------------------------------------------------------------------------- # Initialisation # --------------------------------------------------------------------------- def db_init() -> None: """Create tables and indexes if they do not exist yet.""" if USE_SUPABASE and supabase_client is not None: # Supabase tables are expected to be created via migrations; # we do not attempt DDL here. return conn = _sqlite_conn() try: conn.executescript(""" CREATE TABLE IF NOT EXISTS users ( id TEXT PRIMARY KEY, email TEXT NOT NULL UNIQUE, name TEXT NOT NULL, password TEXT NOT NULL DEFAULT '', plan TEXT NOT NULL DEFAULT 'free', stripe_cid TEXT DEFAULT '', api_key TEXT DEFAULT '', created_at TEXT NOT NULL DEFAULT (datetime('now')) ); CREATE TABLE IF NOT EXISTS sessions ( token TEXT PRIMARY KEY, user_id TEXT NOT NULL, expires_at REAL NOT NULL, FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE ); CREATE TABLE IF NOT EXISTS invoices ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL, vendor TEXT DEFAULT '', inv_number TEXT DEFAULT '', inv_date TEXT DEFAULT '', due_date TEXT DEFAULT '', amount REAL DEFAULT 0, vat_amount REAL DEFAULT 0, total REAL DEFAULT 0, currency TEXT DEFAULT 'USD', status TEXT DEFAULT 'done', is_duplicate INTEGER DEFAULT 0, confidence REAL DEFAULT 0, raw_json TEXT DEFAULT '', created_at TEXT NOT NULL DEFAULT (datetime('now')), FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE ); CREATE INDEX IF NOT EXISTS idx_invoices_user_created ON invoices(user_id, created_at); """) conn.commit() finally: conn.close() # --------------------------------------------------------------------------- # User CRUD # --------------------------------------------------------------------------- def create_user(email: str, name: str, password: str, plan: str = "free") -> dict | None: """Insert a new user. Returns the user dict on success, None on failure.""" uid = str(uuid.uuid4()) if USE_SUPABASE and supabase_client is not None: try: row = ( supabase_client.table("users") .insert({ "id": uid, "email": email, "name": name, "password": password, "plan": plan, "stripe_cid": "", "api_key": "", }) .execute() ) if row.data and len(row.data) > 0: return row.data[0] except Exception: return None return None # SQLite path conn = _sqlite_conn() try: conn.execute( "INSERT INTO users (id, email, name, password, plan) VALUES (?, ?, ?, ?, ?)", (uid, email, name, password, plan), ) conn.commit() cur = conn.execute("SELECT * FROM users WHERE id = ?", (uid,)) return _row_to_dict(cur.fetchone()) except Exception: return None finally: conn.close() def get_user_by_email(email: str) -> dict | None: """Fetch a user by email. Returns user dict or None.""" if USE_SUPABASE and supabase_client is not None: try: row = ( supabase_client.table("users") .select("*") .eq("email", email) .limit(1) .execute() ) if row.data and len(row.data) > 0: return row.data[0] except Exception: pass return None conn = _sqlite_conn() try: cur = conn.execute("SELECT * FROM users WHERE email = ?", (email,)) return _row_to_dict(cur.fetchone()) finally: conn.close() def get_user_by_id(uid: str) -> dict | None: """Fetch a user by primary-key ID. Returns user dict or None.""" if USE_SUPABASE and supabase_client is not None: try: row = ( supabase_client.table("users") .select("*") .eq("id", uid) .limit(1) .execute() ) if row.data and len(row.data) > 0: return row.data[0] except Exception: pass return None conn = _sqlite_conn() try: cur = conn.execute("SELECT * FROM users WHERE id = ?", (uid,)) return _row_to_dict(cur.fetchone()) finally: conn.close() def update_user(uid: str, fields: dict) -> dict | None: """Update arbitrary fields on a user. Returns updated user dict or None.""" if not fields: return get_user_by_id(uid) allowed = {"name", "password", "plan", "stripe_cid", "api_key"} updates = {k: v for k, v in fields.items() if k in allowed} if not updates: return get_user_by_id(uid) if USE_SUPABASE and supabase_client is not None: try: ( supabase_client.table("users") .update(updates) .eq("id", uid) .execute() ) return get_user_by_id(uid) except Exception: return None # SQLite path set_clause = ", ".join(f"{k} = ?" for k in updates) values = list(updates.values()) + [uid] conn = _sqlite_conn() try: conn.execute(f"UPDATE users SET {set_clause} WHERE id = ?", values) conn.commit() return get_user_by_id(uid) except Exception: return None finally: conn.close() # --------------------------------------------------------------------------- # Sessions # --------------------------------------------------------------------------- def create_session(user_id: str, token: str) -> None: """Create a new session row.""" expires_at = time.time() + SESSION_TTL_SECONDS if USE_SUPABASE and supabase_client is not None: try: supabase_client.table("sessions").insert({ "token": token, "user_id": user_id, "expires_at": expires_at, }).execute() except Exception: pass return conn = _sqlite_conn() try: conn.execute( "INSERT INTO sessions (token, user_id, expires_at) VALUES (?, ?, ?)", (token, user_id, expires_at), ) conn.commit() finally: conn.close() def get_session_user(token: str) -> dict | None: """Return the user dict for a valid, non-expired session, or None.""" now = time.time() if USE_SUPABASE and supabase_client is not None: try: row = ( supabase_client.table("sessions") .select("user_id, expires_at") .eq("token", token) .limit(1) .execute() ) if not row.data or len(row.data) == 0: return None session = row.data[0] if session["expires_at"] < now: # Session expired — clean up supabase_client.table("sessions").delete().eq("token", token).execute() return None return get_user_by_id(session["user_id"]) except Exception: return None conn = _sqlite_conn() try: cur = conn.execute( "SELECT user_id, expires_at FROM sessions WHERE token = ?", (token,), ) session = cur.fetchone() if session is None: return None if session["expires_at"] < now: conn.execute("DELETE FROM sessions WHERE token = ?", (token,)) conn.commit() return None return get_user_by_id(session["user_id"]) finally: conn.close() def delete_session(token: str) -> None: """Delete a single session (used on logout).""" if USE_SUPABASE and supabase_client is not None: try: supabase_client.table("sessions").delete().eq("token", token).execute() except Exception: pass return conn = _sqlite_conn() try: conn.execute("DELETE FROM sessions WHERE token = ?", (token,)) conn.commit() finally: conn.close() # --------------------------------------------------------------------------- # Invoices # --------------------------------------------------------------------------- def insert_invoice(user_id: str, data: dict) -> dict | None: """Insert a parsed invoice. *data* should contain the invoice fields. Returns the inserted invoice dict or None on failure. """ inv_id = str(uuid.uuid4()) raw_json = json.dumps(data.get("raw_data", data), default=str) if USE_SUPABASE and supabase_client is not None: try: row = ( supabase_client.table("invoices") .insert({ "id": inv_id, "user_id": user_id, "vendor": data.get("vendor", ""), "inv_number": data.get("inv_number", ""), "inv_date": data.get("inv_date", ""), "due_date": data.get("due_date", ""), "amount": data.get("amount", 0), "vat_amount": data.get("vat_amount", 0), "total": data.get("total", 0), "currency": data.get("currency", "USD"), "status": data.get("status", "done"), "is_duplicate": 1 if data.get("is_duplicate") else 0, "confidence": data.get("confidence", 0), "raw_json": raw_json, }) .execute() ) if row.data and len(row.data) > 0: return row.data[0] except Exception: return None return None # SQLite path conn = _sqlite_conn() try: conn.execute( """INSERT INTO invoices (id, user_id, vendor, inv_number, inv_date, due_date, amount, vat_amount, total, currency, status, is_duplicate, confidence, raw_json) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", ( inv_id, user_id, data.get("vendor", ""), data.get("inv_number", ""), data.get("inv_date", ""), data.get("due_date", ""), data.get("amount", 0), data.get("vat_amount", 0), data.get("total", 0), data.get("currency", "USD"), data.get("status", "done"), 1 if data.get("is_duplicate") else 0, data.get("confidence", 0), raw_json, ), ) conn.commit() cur = conn.execute("SELECT * FROM invoices WHERE id = ?", (inv_id,)) return _row_to_dict(cur.fetchone()) except Exception: return None finally: conn.close() def get_invoices(user_id: str, limit: int = 50, offset: int = 0) -> list[dict]: """Return a list of invoice dicts for *user_id*, newest first.""" if USE_SUPABASE and supabase_client is not None: try: row = ( supabase_client.table("invoices") .select("*") .eq("user_id", user_id) .order("created_at", desc=True) .range(offset, offset + limit - 1) .execute() ) return row.data if row.data else [] except Exception: return [] conn = _sqlite_conn() try: cur = conn.execute( "SELECT * FROM invoices WHERE user_id = ? ORDER BY created_at DESC LIMIT ? OFFSET ?", (user_id, limit, offset), ) return [_row_to_dict(r) for r in cur.fetchall()] finally: conn.close() def count_invoices_this_month(user_id: str) -> int: """Count how many invoices *user_id* created in the current calendar month. Uses an SQL WHERE clause for efficiency (no Python-side filtering). """ import datetime now = datetime.datetime.utcnow() month_start = now.strftime("%Y-%m-01 00:00:00") if USE_SUPABASE and supabase_client is not None: try: row = ( supabase_client.table("invoices") .select("id", count="exact") .eq("user_id", user_id) .gte("created_at", month_start) .execute() ) return row.count if hasattr(row, "count") else len(row.data) except Exception: return 0 conn = _sqlite_conn() try: cur = conn.execute( "SELECT COUNT(*) AS cnt FROM invoices WHERE user_id = ? AND created_at >= ?", (user_id, month_start), ) row = cur.fetchone() return row["cnt"] if row else 0 finally: conn.close() def check_duplicate(user_id: str, vendor: str, total: float) -> bool: """Return True if an invoice with the same vendor and total already exists.""" if USE_SUPABASE and supabase_client is not None: try: row = ( supabase_client.table("invoices") .select("id") .eq("user_id", user_id) .eq("vendor", vendor) .eq("total", total) .limit(1) .execute() ) return bool(row.data and len(row.data) > 0) except Exception: return False conn = _sqlite_conn() try: cur = conn.execute( "SELECT id FROM invoices WHERE user_id = ? AND vendor = ? AND total = ? LIMIT 1", (user_id, vendor, total), ) return cur.fetchone() is not None finally: conn.close() def delete_invoice(user_id: str, invoice_id: str) -> bool: """Delete an invoice belonging to *user_id*. Returns True if a row was removed.""" if USE_SUPABASE and supabase_client is not None: try: row = ( supabase_client.table("invoices") .delete() .eq("id", invoice_id) .eq("user_id", user_id) .execute() ) return bool(row.data and len(row.data) > 0) except Exception: return False conn = _sqlite_conn() try: cur = conn.execute( "DELETE FROM invoices WHERE id = ? AND user_id = ?", (invoice_id, user_id), ) conn.commit() return cur.rowcount > 0 finally: conn.close()