Spaces:
Sleeping
Sleeping
| import os | |
| import sqlite3 | |
| from datetime import datetime, timedelta | |
| from flask import current_app | |
| from utils import get_db_connection | |
| def setup_database(): | |
| """Initializes the database and creates/updates tables as needed.""" | |
| conn = get_db_connection() | |
| cursor = conn.cursor() | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS users ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| username TEXT NOT NULL UNIQUE, | |
| email TEXT NOT NULL UNIQUE, | |
| password_hash TEXT NOT NULL, | |
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP | |
| ); | |
| """) | |
| # Create sessions table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS sessions ( | |
| id TEXT PRIMARY KEY, | |
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
| original_filename TEXT, | |
| persist INTEGER DEFAULT 0, | |
| name TEXT, | |
| user_id INTEGER, | |
| session_type TEXT DEFAULT 'standard' | |
| ); | |
| """) | |
| # Create images table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS images ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| session_id TEXT NOT NULL, | |
| image_index INTEGER NOT NULL, | |
| filename TEXT NOT NULL, | |
| original_name TEXT NOT NULL, | |
| processed_filename TEXT, | |
| image_type TEXT DEFAULT 'original', | |
| box_id TEXT, | |
| FOREIGN KEY (session_id) REFERENCES sessions (id) | |
| ); | |
| """) | |
| # Create questions table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS questions ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| session_id TEXT NOT NULL, | |
| image_id INTEGER NOT NULL, | |
| question_number TEXT, | |
| subject TEXT, | |
| status TEXT, | |
| marked_solution TEXT, | |
| actual_solution TEXT, | |
| time_taken TEXT, | |
| tags TEXT, | |
| FOREIGN KEY (session_id) REFERENCES sessions (id), | |
| FOREIGN KEY (image_id) REFERENCES images (id) | |
| ); | |
| """) | |
| # Create folders table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS folders ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| name TEXT NOT NULL, | |
| parent_id INTEGER, | |
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
| FOREIGN KEY (parent_id) REFERENCES folders (id) ON DELETE CASCADE | |
| ); | |
| """) | |
| # Create generated_pdfs table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS generated_pdfs ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| session_id TEXT NOT NULL, | |
| filename TEXT NOT NULL, | |
| subject TEXT NOT NULL, | |
| tags TEXT, | |
| notes TEXT, | |
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
| source_filename TEXT, | |
| folder_id INTEGER, | |
| persist INTEGER DEFAULT 0, | |
| FOREIGN KEY (session_id) REFERENCES sessions (id), | |
| FOREIGN KEY (folder_id) REFERENCES folders (id) ON DELETE SET NULL | |
| ); | |
| """) | |
| # Create neetprep_questions table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS neetprep_questions ( | |
| id TEXT PRIMARY KEY, | |
| question_text TEXT, | |
| options TEXT, | |
| correct_answer_index INTEGER, | |
| level TEXT, | |
| topic TEXT, | |
| subject TEXT, | |
| last_fetched_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP | |
| ); | |
| """) | |
| # Create neetprep_processed_attempts table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS neetprep_processed_attempts ( | |
| attempt_id TEXT PRIMARY KEY, | |
| processed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP | |
| ); | |
| """) | |
| # Create subjective_folders table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS subjective_folders ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| name TEXT NOT NULL, | |
| parent_id INTEGER, | |
| user_id INTEGER, | |
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
| FOREIGN KEY (parent_id) REFERENCES subjective_folders (id) ON DELETE CASCADE, | |
| FOREIGN KEY (user_id) REFERENCES users (id) | |
| ); | |
| """) | |
| # Create subjective_questions table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS subjective_questions ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| user_id INTEGER NOT NULL, | |
| question_topic TEXT NOT NULL, | |
| question_html TEXT NOT NULL, | |
| question_number_within_topic TEXT, | |
| folder_id INTEGER, | |
| topic_order INTEGER DEFAULT 0, | |
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
| FOREIGN KEY (user_id) REFERENCES users (id), | |
| FOREIGN KEY (folder_id) REFERENCES subjective_folders (id) ON DELETE SET NULL | |
| ); | |
| """) | |
| # Create drive_sources table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS drive_sources ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| name TEXT NOT NULL, | |
| url TEXT NOT NULL, | |
| local_path TEXT, | |
| last_synced TIMESTAMP, | |
| user_id INTEGER, | |
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
| FOREIGN KEY (user_id) REFERENCES users (id) | |
| ); | |
| """) | |
| # Create pdf_access_history table | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS pdf_access_history ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| user_id INTEGER NOT NULL, | |
| file_id TEXT NOT NULL, | |
| filename TEXT NOT NULL, | |
| source_type TEXT DEFAULT 'drive_api', | |
| opened_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
| FOREIGN KEY (user_id) REFERENCES users (id) | |
| ); | |
| """) | |
| # Create qtab_folders table for question table organization | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS qtab_folders ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| name TEXT NOT NULL, | |
| parent_id INTEGER, | |
| user_id INTEGER, | |
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
| FOREIGN KEY (parent_id) REFERENCES qtab_folders (id) ON DELETE CASCADE, | |
| FOREIGN KEY (user_id) REFERENCES users (id) | |
| ); | |
| """) | |
| # Create qtab_images table for question-answer extraction | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS qtab_images ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| user_id INTEGER NOT NULL, | |
| folder_id INTEGER, | |
| filename TEXT NOT NULL, | |
| original_name TEXT NOT NULL, | |
| result_json TEXT, | |
| status TEXT DEFAULT 'pending', | |
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
| FOREIGN KEY (user_id) REFERENCES users (id), | |
| FOREIGN KEY (folder_id) REFERENCES qtab_folders (id) ON DELETE SET NULL | |
| ); | |
| """) | |
| # --- Migrations --- | |
| try: | |
| cursor.execute("SELECT topic_order FROM subjective_questions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE subjective_questions ADD COLUMN topic_order INTEGER DEFAULT 0") | |
| try: | |
| cursor.execute("SELECT folder_id FROM subjective_questions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE subjective_questions ADD COLUMN folder_id INTEGER REFERENCES subjective_folders(id) ON DELETE SET NULL") | |
| try: | |
| cursor.execute("SELECT tags FROM questions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE questions ADD COLUMN tags TEXT") | |
| try: | |
| cursor.execute("SELECT tags FROM questions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE questions ADD COLUMN tags TEXT") | |
| try: | |
| cursor.execute("SELECT image_type FROM images LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE images ADD COLUMN image_type TEXT DEFAULT 'original'") | |
| try: | |
| cursor.execute("SELECT original_filename FROM sessions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE sessions ADD COLUMN original_filename TEXT") | |
| try: | |
| cursor.execute("SELECT persist FROM sessions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE sessions ADD COLUMN persist INTEGER DEFAULT 0") | |
| try: | |
| cursor.execute("SELECT name FROM sessions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE sessions ADD COLUMN name TEXT") | |
| try: | |
| cursor.execute("SELECT persist FROM generated_pdfs LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE generated_pdfs ADD COLUMN persist INTEGER DEFAULT 0") | |
| try: | |
| cursor.execute("SELECT folder_id FROM generated_pdfs LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE generated_pdfs ADD COLUMN folder_id INTEGER REFERENCES folders(id) ON DELETE SET NULL") | |
| try: | |
| cursor.execute("SELECT question_text FROM questions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE questions ADD COLUMN question_text TEXT") | |
| try: | |
| cursor.execute("SELECT chapter FROM questions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE questions ADD COLUMN chapter TEXT") | |
| try: | |
| cursor.execute("SELECT subject FROM sessions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE sessions ADD COLUMN subject TEXT") | |
| try: | |
| cursor.execute("SELECT tags FROM sessions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE sessions ADD COLUMN tags TEXT") | |
| try: | |
| cursor.execute("SELECT notes FROM sessions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE sessions ADD COLUMN notes TEXT") | |
| # --- Multi-user Migrations --- | |
| try: | |
| cursor.execute("SELECT user_id FROM sessions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE sessions ADD COLUMN user_id INTEGER REFERENCES users(id)") | |
| try: | |
| cursor.execute("SELECT user_id FROM generated_pdfs LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE generated_pdfs ADD COLUMN user_id INTEGER REFERENCES users(id)") | |
| try: | |
| cursor.execute("SELECT user_id FROM folders LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE folders ADD COLUMN user_id INTEGER REFERENCES users(id)") | |
| try: | |
| cursor.execute("SELECT neetprep_enabled FROM users LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE users ADD COLUMN neetprep_enabled INTEGER DEFAULT 1") | |
| try: | |
| cursor.execute("SELECT dpi FROM users LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE users ADD COLUMN dpi INTEGER DEFAULT 100") | |
| try: | |
| cursor.execute("SELECT color_rm_dpi FROM users LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE users ADD COLUMN color_rm_dpi INTEGER DEFAULT 200") | |
| try: | |
| cursor.execute("SELECT box_id FROM images LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE images ADD COLUMN box_id TEXT") | |
| try: | |
| cursor.execute("SELECT session_type FROM sessions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE sessions ADD COLUMN session_type TEXT DEFAULT 'standard'") | |
| try: | |
| cursor.execute("SELECT v2_default FROM users LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE users ADD COLUMN v2_default INTEGER DEFAULT 0") | |
| try: | |
| cursor.execute("SELECT magnifier_enabled FROM users LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE users ADD COLUMN magnifier_enabled INTEGER DEFAULT 1") | |
| try: | |
| cursor.execute("SELECT source_type FROM drive_sources LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE drive_sources ADD COLUMN source_type TEXT DEFAULT 'folder'") | |
| try: | |
| cursor.execute("SELECT google_token FROM users LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE users ADD COLUMN google_token TEXT") | |
| try: | |
| cursor.execute("SELECT question_json FROM subjective_questions LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE subjective_questions ADD COLUMN question_json TEXT") | |
| try: | |
| cursor.execute("SELECT classifier_model FROM users LIMIT 1") | |
| except sqlite3.OperationalError: | |
| cursor.execute("ALTER TABLE users ADD COLUMN classifier_model TEXT DEFAULT 'gemini'") | |
| conn.commit() | |
| conn.close() | |
| def cleanup_old_data(): | |
| """Removes sessions, files, and PDFs older than 1 day, unless persisted.""" | |
| print("Running cleanup of old data...") | |
| conn = get_db_connection() | |
| cutoff = datetime.now() - timedelta(days=1) | |
| old_sessions = conn.execute('SELECT id FROM sessions WHERE created_at < ? AND persist = 0', (cutoff,)).fetchall() | |
| for session in old_sessions: | |
| session_id = session['id'] | |
| print(f"Deleting old session: {session_id}") | |
| images_to_delete = conn.execute('SELECT filename, processed_filename FROM images WHERE session_id = ?', (session_id,)).fetchall() | |
| for img in images_to_delete: | |
| if img['filename']: | |
| try: os.remove(os.path.join(current_app.config['UPLOAD_FOLDER'], img['filename'])) | |
| except OSError: pass | |
| if img['processed_filename']: | |
| try: os.remove(os.path.join(current_app.config['PROCESSED_FOLDER'], img['processed_filename'])) | |
| except OSError: pass | |
| conn.execute('DELETE FROM questions WHERE session_id = ?', (session_id,)) | |
| conn.execute('DELETE FROM images WHERE session_id = ?', (session_id,)) | |
| conn.execute('DELETE FROM sessions WHERE id = ?', (session_id,)) | |
| old_pdfs = conn.execute('SELECT id, filename FROM generated_pdfs WHERE created_at < ? AND persist = 0', (cutoff,)).fetchall() | |
| for pdf in old_pdfs: | |
| pdf_id, pdf_filename = pdf['id'], pdf['filename'] | |
| print(f"Deleting old generated PDF: {pdf_filename}") | |
| try: | |
| os.remove(os.path.join(current_app.config['OUTPUT_FOLDER'], pdf_filename)) | |
| except OSError: | |
| pass | |
| conn.execute('DELETE FROM generated_pdfs WHERE id = ?', (pdf_id,)) | |
| db_filenames = {row['filename'] for row in conn.execute('SELECT filename FROM generated_pdfs').fetchall()} | |
| for filename in os.listdir(current_app.config['OUTPUT_FOLDER']): | |
| if filename not in db_filenames: | |
| file_path = os.path.join(current_app.config['OUTPUT_FOLDER'], filename) | |
| file_mtime = datetime.fromtimestamp(os.path.getmtime(file_path)) | |
| if file_mtime < cutoff: | |
| print(f"Deleting old, orphaned PDF: {filename}") | |
| try: | |
| os.remove(file_path) | |
| except OSError: | |
| pass | |
| conn.commit() | |
| conn.close() | |
| print("Cleanup finished.") | |
| def get_folder_tree(user_id=None): | |
| conn = get_db_connection() | |
| if user_id: | |
| folders = conn.execute('SELECT id, name, parent_id FROM folders WHERE user_id = ? ORDER BY name', (user_id,)).fetchall() | |
| else: | |
| # Fallback for old behavior or admin views | |
| folders = conn.execute('SELECT id, name, parent_id FROM folders ORDER BY name').fetchall() | |
| conn.close() | |
| folder_map = {f['id']: dict(f) for f in folders} | |
| tree = [] | |
| for folder_id, folder in folder_map.items(): | |
| if folder['parent_id']: | |
| parent = folder_map.get(folder['parent_id']) | |
| if parent: | |
| if 'children' not in parent: | |
| parent['children'] = [] | |
| parent['children'].append(folder) | |
| else: | |
| tree.append(folder) | |
| return tree | |
| def get_subjective_folder_tree(user_id): | |
| conn = get_db_connection() | |
| folders = conn.execute('SELECT id, name, parent_id FROM subjective_folders WHERE user_id = ? ORDER BY name', (user_id,)).fetchall() | |
| conn.close() | |
| folder_map = {f['id']: dict(f) for f in folders} | |
| tree = [] | |
| for folder_id, folder in folder_map.items(): | |
| if folder['parent_id']: | |
| parent = folder_map.get(folder['parent_id']) | |
| if parent: | |
| if 'children' not in parent: | |
| parent['children'] = [] | |
| parent['children'].append(folder) | |
| else: | |
| tree.append(folder) | |
| return tree | |
| def get_qtab_folder_tree(user_id): | |
| conn = get_db_connection() | |
| folders = conn.execute('SELECT id, name, parent_id FROM qtab_folders WHERE user_id = ? ORDER BY name', (user_id,)).fetchall() | |
| conn.close() | |
| folder_map = {f['id']: dict(f) for f in folders} | |
| tree = [] | |
| for folder_id, folder in folder_map.items(): | |
| if folder['parent_id']: | |
| parent = folder_map.get(folder['parent_id']) | |
| if parent: | |
| if 'children' not in parent: | |
| parent['children'] = [] | |
| parent['children'].append(folder) | |
| else: | |
| tree.append(folder) | |
| return tree | |
| def get_all_descendant_folder_ids(conn, folder_id, user_id=None): | |
| """Recursively gets all descendant folder IDs for a given folder, scoped to a user.""" | |
| if user_id: | |
| children = conn.execute('SELECT id FROM folders WHERE parent_id = ? AND user_id = ?', (folder_id, user_id)).fetchall() | |
| else: | |
| children = conn.execute('SELECT id FROM folders WHERE parent_id = ?', (folder_id,)).fetchall() | |
| folder_ids = [f['id'] for f in children] | |
| for child_id in folder_ids: | |
| folder_ids.extend(get_all_descendant_folder_ids(conn, child_id, user_id)) | |
| return folder_ids | |