Spaces:
Sleeping
Sleeping
| """ | |
| Database layer — SQLite for user data and thread metadata. | |
| LangGraph checkpointing also uses SQLite (configured in graph.py). | |
| """ | |
| import sqlite3 | |
| import time | |
| from config import DB_PATH | |
| from security import encrypt_secret, decrypt_secret | |
| # Autocommit mode (isolation_level=None) makes every write atomic and turns | |
| # conn.commit() into a harmless no-op — this avoids the "cannot commit - no | |
| # transaction is active" race that occurred when the single shared connection | |
| # was used concurrently by FastAPI's worker threads (e.g. a double-fired DELETE). | |
| # WAL allows concurrent readers; busy_timeout waits instead of erroring on locks. | |
| conn = sqlite3.connect(database=DB_PATH, check_same_thread=False, isolation_level=None) | |
| conn.execute("PRAGMA journal_mode=WAL") | |
| conn.execute("PRAGMA busy_timeout=5000") | |
| conn.execute("PRAGMA synchronous=NORMAL") | |
| # --- Schema --- | |
| conn.execute(""" | |
| CREATE TABLE IF NOT EXISTS users ( | |
| google_id TEXT PRIMARY KEY, | |
| email TEXT NOT NULL, | |
| name TEXT NOT NULL, | |
| picture TEXT, | |
| openrouter_key TEXT, | |
| tavily_key TEXT, | |
| student_profile TEXT, | |
| created_at REAL NOT NULL, | |
| last_login REAL NOT NULL | |
| ) | |
| """) | |
| # Upgrade path: add tavily_key column to pre-existing databases. | |
| try: | |
| conn.execute("ALTER TABLE users ADD COLUMN tavily_key TEXT") | |
| except sqlite3.OperationalError: | |
| pass # column already exists | |
| conn.execute(""" | |
| CREATE TABLE IF NOT EXISTS thread_titles ( | |
| thread_id TEXT PRIMARY KEY, | |
| title TEXT, | |
| last_active REAL, | |
| user_id TEXT | |
| ) | |
| """) | |
| conn.execute(""" | |
| CREATE TABLE IF NOT EXISTS notes ( | |
| note_id TEXT PRIMARY KEY, | |
| user_id TEXT NOT NULL, | |
| title TEXT, | |
| content TEXT, | |
| created_at REAL, | |
| updated_at REAL | |
| ) | |
| """) | |
| conn.commit() | |
| # --- User operations --- | |
| def upsert_user(google_id: str, email: str, name: str, picture: str = "") -> dict: | |
| """Create or update a user from Google OAuth data.""" | |
| now = time.time() | |
| existing = conn.execute( | |
| "SELECT google_id FROM users WHERE google_id = ?", (google_id,) | |
| ).fetchone() | |
| if existing: | |
| conn.execute( | |
| "UPDATE users SET last_login = ?, name = ?, picture = ? WHERE google_id = ?", | |
| (now, name, picture, google_id), | |
| ) | |
| else: | |
| conn.execute( | |
| "INSERT INTO users (google_id, email, name, picture, created_at, last_login) " | |
| "VALUES (?, ?, ?, ?, ?, ?)", | |
| (google_id, email, name, picture, now, now), | |
| ) | |
| conn.commit() | |
| return {"google_id": google_id, "email": email, "name": name, "picture": picture} | |
| def get_user(google_id: str) -> dict | None: | |
| row = conn.execute( | |
| "SELECT google_id, email, name, picture, openrouter_key, student_profile, tavily_key " | |
| "FROM users WHERE google_id = ?", | |
| (google_id,), | |
| ).fetchone() | |
| if not row: | |
| return None | |
| return { | |
| "google_id": row[0], "email": row[1], "name": row[2], | |
| "picture": row[3], "openrouter_key": decrypt_secret(row[4] or ""), | |
| "student_profile": row[5] or "", | |
| "tavily_key": decrypt_secret(row[6] or ""), | |
| } | |
| def save_user_api_key(google_id: str, key: str): | |
| conn.execute( | |
| "UPDATE users SET openrouter_key = ? WHERE google_id = ?", | |
| (encrypt_secret(key), google_id), | |
| ) | |
| conn.commit() | |
| def save_tavily_key(google_id: str, key: str): | |
| conn.execute( | |
| "UPDATE users SET tavily_key = ? WHERE google_id = ?", | |
| (encrypt_secret(key), google_id), | |
| ) | |
| conn.commit() | |
| def get_tavily_key(google_id: str) -> str: | |
| row = conn.execute( | |
| "SELECT tavily_key FROM users WHERE google_id = ?", (google_id,) | |
| ).fetchone() | |
| return decrypt_secret(row[0]) if row and row[0] else "" | |
| def save_student_profile(google_id: str, profile: str): | |
| conn.execute( | |
| "UPDATE users SET student_profile = ? WHERE google_id = ?", (profile, google_id) | |
| ) | |
| conn.commit() | |
| def get_user_api_key(google_id: str) -> str: | |
| row = conn.execute( | |
| "SELECT openrouter_key FROM users WHERE google_id = ?", (google_id,) | |
| ).fetchone() | |
| return decrypt_secret(row[0]) if row and row[0] else "" | |
| # --- Thread operations --- | |
| def get_threads(user_id: str = ""): | |
| if user_id: | |
| cursor = conn.execute( | |
| "SELECT thread_id, title FROM thread_titles " | |
| "WHERE user_id = ? ORDER BY last_active DESC", | |
| (user_id,), | |
| ) | |
| else: | |
| cursor = conn.execute( | |
| "SELECT thread_id, title FROM thread_titles ORDER BY last_active DESC" | |
| ) | |
| return [{"id": row[0], "title": row[1]} for row in cursor] | |
| def upsert_thread(thread_id: str, message: str, now: float, user_id: str = ""): | |
| row = conn.execute( | |
| "SELECT thread_id FROM thread_titles WHERE thread_id = ?", (thread_id,) | |
| ).fetchone() | |
| if row: | |
| conn.execute( | |
| "UPDATE thread_titles SET last_active = ? WHERE thread_id = ?", | |
| (now, thread_id), | |
| ) | |
| else: | |
| title = message[:30] + ("..." if len(message) > 30 else "") | |
| conn.execute( | |
| "INSERT INTO thread_titles (thread_id, title, last_active, user_id) " | |
| "VALUES (?, ?, ?, ?)", | |
| (thread_id, title, now, user_id), | |
| ) | |
| conn.commit() | |
| def get_thread_owner(thread_id: str) -> str | None: | |
| """Return the user_id that owns a thread, or None if the thread is unknown.""" | |
| row = conn.execute( | |
| "SELECT user_id FROM thread_titles WHERE thread_id = ?", (thread_id,) | |
| ).fetchone() | |
| return (row[0] or None) if row else None | |
| def rename_thread(thread_id: str, title: str): | |
| conn.execute( | |
| "UPDATE thread_titles SET title = ? WHERE thread_id = ?", (title, thread_id) | |
| ) | |
| conn.commit() | |
| def delete_thread(thread_id: str): | |
| conn.execute("DELETE FROM thread_titles WHERE thread_id = ?", (thread_id,)) | |
| conn.commit() | |
| # --- Notes (Canvas) operations --- | |
| def list_notes(user_id: str): | |
| """List a user's notes (metadata only, newest first) for the Canvas list view.""" | |
| cursor = conn.execute( | |
| "SELECT note_id, title, updated_at FROM notes " | |
| "WHERE user_id = ? ORDER BY updated_at DESC", | |
| (user_id,), | |
| ) | |
| return [ | |
| {"note_id": row[0], "title": row[1], "updated_at": row[2]} | |
| for row in cursor | |
| ] | |
| def get_note_owner(note_id: str) -> str | None: | |
| """Return the user_id that owns a note, or None if the note is unknown.""" | |
| row = conn.execute( | |
| "SELECT user_id FROM notes WHERE note_id = ?", (note_id,) | |
| ).fetchone() | |
| return (row[0] or None) if row else None | |
| def get_note(note_id: str) -> dict | None: | |
| row = conn.execute( | |
| "SELECT note_id, user_id, title, content, created_at, updated_at " | |
| "FROM notes WHERE note_id = ?", | |
| (note_id,), | |
| ).fetchone() | |
| if not row: | |
| return None | |
| return { | |
| "note_id": row[0], "user_id": row[1], "title": row[2], | |
| "content": row[3] or "", "created_at": row[4], "updated_at": row[5], | |
| } | |
| def create_note(note_id: str, user_id: str, title: str, content: str, now: float): | |
| conn.execute( | |
| "INSERT INTO notes (note_id, user_id, title, content, created_at, updated_at) " | |
| "VALUES (?, ?, ?, ?, ?, ?)", | |
| (note_id, user_id, title, content, now, now), | |
| ) | |
| conn.commit() | |
| def update_note(note_id: str, title: str, content: str, now: float): | |
| conn.execute( | |
| "UPDATE notes SET title = ?, content = ?, updated_at = ? WHERE note_id = ?", | |
| (title, content, now, note_id), | |
| ) | |
| conn.commit() | |
| def delete_note(note_id: str): | |
| conn.execute("DELETE FROM notes WHERE note_id = ?", (note_id,)) | |
| conn.commit() | |