Spaces:
Paused
Paused
| """SQLite database for ClassLens - users, sessions, parsed data, prompts, reports.""" | |
| from __future__ import annotations | |
| import aiosqlite | |
| import json | |
| from pathlib import Path | |
| from datetime import datetime | |
| from typing import Optional | |
| from .config import get_settings | |
| DATABASE_PATH = Path(__file__).parent.parent / "classlens.db" | |
| def _get_db_path() -> Path: | |
| settings = get_settings() | |
| if settings.database_path: | |
| return Path(settings.database_path) | |
| return DATABASE_PATH | |
| async def init_database(): | |
| """Initialize the database with required tables.""" | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| await db.execute(""" | |
| CREATE TABLE IF NOT EXISTS users ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| email TEXT UNIQUE NOT NULL, | |
| password_hash TEXT NOT NULL, | |
| display_name TEXT, | |
| created_at TEXT NOT NULL, | |
| last_login TEXT | |
| ) | |
| """) | |
| await db.execute(""" | |
| CREATE TABLE IF NOT EXISTS exam_sessions ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| user_id INTEGER NOT NULL, | |
| title TEXT DEFAULT 'Untitled Session', | |
| status TEXT DEFAULT 'draft', | |
| created_at TEXT NOT NULL, | |
| updated_at TEXT NOT NULL, | |
| FOREIGN KEY (user_id) REFERENCES users(id) | |
| ) | |
| """) | |
| await db.execute(""" | |
| CREATE TABLE IF NOT EXISTS parsed_data ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| session_id INTEGER NOT NULL, | |
| data_type TEXT NOT NULL, | |
| file_name TEXT, | |
| raw_text TEXT, | |
| structured_data TEXT NOT NULL, | |
| created_at TEXT NOT NULL, | |
| FOREIGN KEY (session_id) REFERENCES exam_sessions(id) ON DELETE CASCADE | |
| ) | |
| """) | |
| await db.execute(""" | |
| CREATE TABLE IF NOT EXISTS prompts ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| user_id INTEGER NOT NULL, | |
| name TEXT NOT NULL, | |
| content TEXT NOT NULL, | |
| is_default BOOLEAN DEFAULT 0, | |
| created_at TEXT NOT NULL, | |
| updated_at TEXT NOT NULL, | |
| FOREIGN KEY (user_id) REFERENCES users(id) | |
| ) | |
| """) | |
| await db.execute(""" | |
| CREATE TABLE IF NOT EXISTS reports ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| session_id INTEGER NOT NULL, | |
| prompt_id INTEGER, | |
| html_content TEXT NOT NULL, | |
| created_at TEXT NOT NULL, | |
| FOREIGN KEY (session_id) REFERENCES exam_sessions(id) ON DELETE CASCADE, | |
| FOREIGN KEY (prompt_id) REFERENCES prompts(id) | |
| ) | |
| """) | |
| await db.commit() | |
| # ---- User CRUD ---- | |
| async def create_user(email: str, password_hash: str, display_name: str = "") -> int: | |
| db_path = _get_db_path() | |
| now = datetime.utcnow().isoformat() | |
| async with aiosqlite.connect(db_path) as db: | |
| cursor = await db.execute( | |
| "INSERT INTO users (email, password_hash, display_name, created_at) VALUES (?, ?, ?, ?)", | |
| (email, password_hash, display_name or email.split("@")[0], now), | |
| ) | |
| await db.commit() | |
| return cursor.lastrowid | |
| async def get_user_by_email(email: str) -> Optional[dict]: | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute("SELECT * FROM users WHERE email = ?", (email,)) as cursor: | |
| row = await cursor.fetchone() | |
| return dict(row) if row else None | |
| async def get_user_by_id(user_id: int) -> Optional[dict]: | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute("SELECT * FROM users WHERE id = ?", (user_id,)) as cursor: | |
| row = await cursor.fetchone() | |
| return dict(row) if row else None | |
| async def update_last_login(user_id: int): | |
| db_path = _get_db_path() | |
| now = datetime.utcnow().isoformat() | |
| async with aiosqlite.connect(db_path) as db: | |
| await db.execute("UPDATE users SET last_login = ? WHERE id = ?", (now, user_id)) | |
| await db.commit() | |
| # ---- Session CRUD ---- | |
| async def create_session(user_id: int, title: str = "Untitled Session") -> int: | |
| db_path = _get_db_path() | |
| now = datetime.utcnow().isoformat() | |
| async with aiosqlite.connect(db_path) as db: | |
| cursor = await db.execute( | |
| "INSERT INTO exam_sessions (user_id, title, status, created_at, updated_at) VALUES (?, ?, 'draft', ?, ?)", | |
| (user_id, title, now, now), | |
| ) | |
| await db.commit() | |
| return cursor.lastrowid | |
| async def get_sessions(user_id: int) -> list[dict]: | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute( | |
| "SELECT * FROM exam_sessions WHERE user_id = ? ORDER BY created_at DESC", (user_id,) | |
| ) as cursor: | |
| rows = await cursor.fetchall() | |
| return [dict(row) for row in rows] | |
| async def get_session(session_id: int) -> Optional[dict]: | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute("SELECT * FROM exam_sessions WHERE id = ?", (session_id,)) as cursor: | |
| row = await cursor.fetchone() | |
| return dict(row) if row else None | |
| async def update_session_status(session_id: int, status: str): | |
| db_path = _get_db_path() | |
| now = datetime.utcnow().isoformat() | |
| async with aiosqlite.connect(db_path) as db: | |
| await db.execute( | |
| "UPDATE exam_sessions SET status = ?, updated_at = ? WHERE id = ?", | |
| (status, now, session_id), | |
| ) | |
| await db.commit() | |
| # ---- Parsed Data CRUD ---- | |
| async def save_parsed_data( | |
| session_id: int, data_type: str, file_name: str, raw_text: str, structured_data: dict | |
| ) -> int: | |
| db_path = _get_db_path() | |
| now = datetime.utcnow().isoformat() | |
| async with aiosqlite.connect(db_path) as db: | |
| cursor = await db.execute( | |
| "INSERT INTO parsed_data (session_id, data_type, file_name, raw_text, structured_data, created_at) VALUES (?, ?, ?, ?, ?, ?)", | |
| (session_id, data_type, file_name, raw_text, json.dumps(structured_data, ensure_ascii=False), now), | |
| ) | |
| await db.commit() | |
| return cursor.lastrowid | |
| async def get_parsed_data(session_id: int) -> list[dict]: | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute( | |
| "SELECT * FROM parsed_data WHERE session_id = ? ORDER BY data_type, created_at", | |
| (session_id,), | |
| ) as cursor: | |
| rows = await cursor.fetchall() | |
| result = [] | |
| for row in rows: | |
| d = dict(row) | |
| d["structured_data"] = json.loads(d["structured_data"]) | |
| result.append(d) | |
| return result | |
| async def delete_parsed_data(session_id: int, data_type: Optional[str] = None): | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| if data_type: | |
| await db.execute( | |
| "DELETE FROM parsed_data WHERE session_id = ? AND data_type = ?", | |
| (session_id, data_type), | |
| ) | |
| else: | |
| await db.execute("DELETE FROM parsed_data WHERE session_id = ?", (session_id,)) | |
| await db.commit() | |
| # ---- Prompt CRUD ---- | |
| async def save_prompt(user_id: int, name: str, content: str, is_default: bool = False) -> int: | |
| db_path = _get_db_path() | |
| now = datetime.utcnow().isoformat() | |
| async with aiosqlite.connect(db_path) as db: | |
| cursor = await db.execute( | |
| "INSERT INTO prompts (user_id, name, content, is_default, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", | |
| (user_id, name, content, is_default, now, now), | |
| ) | |
| await db.commit() | |
| return cursor.lastrowid | |
| async def get_prompts(user_id: int) -> list[dict]: | |
| """Get user's own prompts plus all default prompts.""" | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute( | |
| "SELECT * FROM prompts WHERE user_id = ? OR is_default = 1 ORDER BY is_default DESC, updated_at DESC", | |
| (user_id,), | |
| ) as cursor: | |
| rows = await cursor.fetchall() | |
| return [dict(row) for row in rows] | |
| async def get_all_prompts() -> list[dict]: | |
| """Get all prompts from all users (for the 'read others' dropdown).""" | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute( | |
| "SELECT p.*, u.email as author_email FROM prompts p JOIN users u ON p.user_id = u.id ORDER BY p.updated_at DESC" | |
| ) as cursor: | |
| rows = await cursor.fetchall() | |
| return [dict(row) for row in rows] | |
| async def get_prompt(prompt_id: int) -> Optional[dict]: | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute("SELECT * FROM prompts WHERE id = ?", (prompt_id,)) as cursor: | |
| row = await cursor.fetchone() | |
| return dict(row) if row else None | |
| async def update_prompt(prompt_id: int, name: str, content: str): | |
| db_path = _get_db_path() | |
| now = datetime.utcnow().isoformat() | |
| async with aiosqlite.connect(db_path) as db: | |
| await db.execute( | |
| "UPDATE prompts SET name = ?, content = ?, updated_at = ? WHERE id = ?", | |
| (name, content, now, prompt_id), | |
| ) | |
| await db.commit() | |
| async def delete_prompt(prompt_id: int): | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| await db.execute("DELETE FROM prompts WHERE id = ?", (prompt_id,)) | |
| await db.commit() | |
| # ---- Report CRUD ---- | |
| async def save_report(session_id: int, prompt_id: Optional[int], html_content: str) -> int: | |
| db_path = _get_db_path() | |
| now = datetime.utcnow().isoformat() | |
| async with aiosqlite.connect(db_path) as db: | |
| cursor = await db.execute( | |
| "INSERT INTO reports (session_id, prompt_id, html_content, created_at) VALUES (?, ?, ?, ?)", | |
| (session_id, prompt_id, html_content, now), | |
| ) | |
| await db.commit() | |
| return cursor.lastrowid | |
| async def get_reports(session_id: int) -> list[dict]: | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute( | |
| "SELECT * FROM reports WHERE session_id = ? ORDER BY created_at DESC", | |
| (session_id,), | |
| ) as cursor: | |
| rows = await cursor.fetchall() | |
| return [dict(row) for row in rows] | |
| async def get_report(report_id: int) -> Optional[dict]: | |
| db_path = _get_db_path() | |
| async with aiosqlite.connect(db_path) as db: | |
| db.row_factory = aiosqlite.Row | |
| async with db.execute("SELECT * FROM reports WHERE id = ?", (report_id,)) as cursor: | |
| row = await cursor.fetchone() | |
| return dict(row) if row else None | |