| """ |
| 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 |
|
|
| |
| |
| |
|
|
| _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) |
|
|
|
|
| |
| |
| |
|
|
| def db_init() -> None: |
| """Create tables and indexes if they do not exist yet.""" |
| if USE_SUPABASE and supabase_client is not None: |
| |
| |
| 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() |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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() |
|
|
|
|
| |
| |
| |
|
|
| 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: |
| |
| 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() |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
| |
| 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() |
