Spaces:
Sleeping
Sleeping
| """SQLite database for generation history.""" | |
| from __future__ import annotations | |
| import json | |
| import sqlite3 | |
| import logging | |
| from datetime import datetime, timedelta, timezone | |
| from pathlib import Path | |
| logger = logging.getLogger(__name__) | |
| _SCHEMA = """ | |
| CREATE TABLE IF NOT EXISTS generation_history ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| job_id TEXT NOT NULL UNIQUE, | |
| session_id TEXT NOT NULL, | |
| original_filename TEXT NOT NULL, | |
| original_path TEXT NOT NULL, | |
| output_path TEXT NOT NULL, | |
| difficulty TEXT NOT NULL, | |
| num_differences INTEGER NOT NULL, | |
| processing_time REAL NOT NULL, | |
| metadata TEXT NOT NULL, | |
| created_at TEXT NOT NULL, | |
| expires_at TEXT NOT NULL | |
| ); | |
| CREATE INDEX IF NOT EXISTS idx_job_id ON generation_history(job_id); | |
| CREATE INDEX IF NOT EXISTS idx_expires_at ON generation_history(expires_at); | |
| CREATE TABLE IF NOT EXISTS job_status ( | |
| job_id TEXT PRIMARY KEY, | |
| status TEXT NOT NULL, | |
| progress INTEGER NOT NULL DEFAULT 0, | |
| current_step TEXT NOT NULL DEFAULT '', | |
| error TEXT, | |
| result_path TEXT, | |
| created_at TEXT NOT NULL, | |
| updated_at TEXT NOT NULL | |
| ); | |
| CREATE INDEX IF NOT EXISTS idx_job_status_updated ON job_status(updated_at); | |
| """ | |
| def init_db(db_path: str) -> None: | |
| """Create the database and tables if they don't exist. | |
| Raises OSError if the database directory cannot be created. | |
| """ | |
| db_file = Path(db_path) | |
| try: | |
| # Create parent directory if needed | |
| db_file.parent.mkdir(parents=True, exist_ok=True) | |
| logger.info(f"Database directory ensured: {db_file.parent}") | |
| except OSError as e: | |
| logger.error(f"Failed to create database directory {db_file.parent}: {e}") | |
| raise OSError( | |
| f"Cannot create database directory {db_file.parent}. " | |
| f"Please ensure FLASK_ENV=production is set for deployment." | |
| ) from e | |
| try: | |
| # Create database and schema | |
| with sqlite3.connect(db_path) as conn: | |
| conn.executescript(_SCHEMA) | |
| logger.info(f"Database initialized: {db_path}") | |
| except sqlite3.Error as e: | |
| logger.error(f"Failed to initialize database: {e}") | |
| raise | |
| def save_generation( | |
| db_path: str, | |
| job_id: str, | |
| session_id: str, | |
| original_filename: str, | |
| original_path: str, | |
| output_path: str, | |
| difficulty: str, | |
| num_differences: int, | |
| processing_time: float, | |
| metadata: dict, | |
| expiry_hours: int = 24, | |
| ) -> int: | |
| """Save a generation record. Returns the row id.""" | |
| now = datetime.now(timezone.utc) | |
| expires = now + timedelta(hours=expiry_hours) | |
| with sqlite3.connect(db_path) as conn: | |
| cursor = conn.execute( | |
| """INSERT INTO generation_history | |
| (job_id, session_id, original_filename, original_path, | |
| output_path, difficulty, num_differences, processing_time, | |
| metadata, created_at, expires_at) | |
| VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", | |
| ( | |
| job_id, | |
| session_id, | |
| original_filename, | |
| original_path, | |
| output_path, | |
| difficulty, | |
| num_differences, | |
| processing_time, | |
| json.dumps(metadata, ensure_ascii=False), | |
| now.isoformat(), | |
| expires.isoformat(), | |
| ), | |
| ) | |
| return cursor.lastrowid | |
| def get_generation(db_path: str, job_id: str) -> dict | None: | |
| """Retrieve a generation record by job_id.""" | |
| with sqlite3.connect(db_path) as conn: | |
| conn.row_factory = sqlite3.Row | |
| row = conn.execute( | |
| "SELECT * FROM generation_history WHERE job_id = ?", (job_id,) | |
| ).fetchone() | |
| if row is None: | |
| return None | |
| d = dict(row) | |
| d["metadata"] = json.loads(d["metadata"]) | |
| return d | |
| def cleanup_expired(db_path: str) -> int: | |
| """Delete expired records. Returns count of deleted rows.""" | |
| now = datetime.now(timezone.utc).isoformat() | |
| with sqlite3.connect(db_path) as conn: | |
| cursor = conn.execute( | |
| "DELETE FROM generation_history WHERE expires_at < ?", (now,) | |
| ) | |
| return cursor.rowcount | |
| # Job status persistence functions for JobManager | |
| def save_job_status( | |
| db_path: str, | |
| job_id: str, | |
| status: str, | |
| progress: int = 0, | |
| current_step: str = "", | |
| error: str | None = None, | |
| result_path: str | None = None, | |
| ) -> None: | |
| """Save or update job status in database.""" | |
| now = datetime.now(timezone.utc).isoformat() | |
| with sqlite3.connect(db_path) as conn: | |
| # Use INSERT OR REPLACE to handle both new and existing jobs | |
| conn.execute( | |
| """INSERT OR REPLACE INTO job_status | |
| (job_id, status, progress, current_step, error, result_path, | |
| created_at, updated_at) | |
| VALUES ( | |
| ?, | |
| ?, | |
| ?, | |
| ?, | |
| ?, | |
| ?, | |
| COALESCE((SELECT created_at FROM job_status WHERE job_id = ?), ?), | |
| ? | |
| )""", | |
| ( | |
| job_id, | |
| status, | |
| progress, | |
| current_step, | |
| error, | |
| result_path, | |
| job_id, | |
| now, | |
| now, | |
| ), | |
| ) | |
| def get_job_status(db_path: str, job_id: str) -> dict | None: | |
| """Retrieve job status by job_id.""" | |
| with sqlite3.connect(db_path) as conn: | |
| conn.row_factory = sqlite3.Row | |
| row = conn.execute( | |
| "SELECT * FROM job_status WHERE job_id = ?", (job_id,) | |
| ).fetchone() | |
| if row is None: | |
| return None | |
| return dict(row) | |
| def cleanup_old_job_status(db_path: str, hours: int = 24) -> int: | |
| """Delete job status records older than specified hours. Returns count of deleted rows.""" | |
| cutoff = (datetime.now(timezone.utc) - timedelta(hours=hours)).isoformat() | |
| with sqlite3.connect(db_path) as conn: | |
| cursor = conn.execute( | |
| "DELETE FROM job_status WHERE updated_at < ?", (cutoff,) | |
| ) | |
| return cursor.rowcount | |