import aiosqlite import logging from typing import Optional, List, Dict, Any import json import os logger = logging.getLogger(__name__) # Get the absolute path to studspider.db in the project root DEFAULT_DB_PATH = os.path.join( os.path.dirname(os.path.dirname(__file__)), "studspider.db" ) if not os.path.exists(DEFAULT_DB_PATH): raise FileNotFoundError(f"Database file not found at {DEFAULT_DB_PATH}") print('Please run "src/scripts/studspider.SQL" to setup the database.') os.exit(0) class AsyncSchedulerDB: """ Asynchronous SQLite database manager for APScheduler. Handles job storage, retrieval, and database operations using the main studspider.db database. """ def __init__(self, db_path: str = DEFAULT_DB_PATH): """ Initialize the database manager. Args: db_path (str, optional): Path to the SQLite database file. Defaults to studspider.db in project root. """ self.db_path = db_path self._initialized = False logger.info(f"Using database at: {self.db_path}") async def initialize(self): """Initialize the database schema if it doesn't exist.""" if self._initialized: return async with aiosqlite.connect(self.db_path) as db: # Create apscheduler jobs table if it doesn't exist await db.execute( """ CREATE TABLE IF NOT EXISTS apscheduler_jobs ( id VARCHAR(191) NOT NULL, next_run_time FLOAT, job_state BLOB NOT NULL, PRIMARY KEY (id) ) """ ) await db.commit() self._initialized = True logger.info(f"Initialized scheduler database at {self.db_path}") async def get_all_jobs(self) -> List[Dict[str, Any]]: """ Retrieve all jobs from the database. Returns: List[Dict[str, Any]]: List of job dictionaries """ await self.initialize() async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM apscheduler_jobs ORDER BY next_run_time" ) as cursor: rows = await cursor.fetchall() return [ { "id": row["id"], "next_run_time": row["next_run_time"], "job_state": json.loads(row["job_state"]), } for row in rows ] async def get_job(self, job_id: str) -> Optional[Dict[str, Any]]: """ Retrieve a specific job by ID. Args: job_id (str): The ID of the job to retrieve Returns: Optional[Dict[str, Any]]: Job dictionary if found, None otherwise """ await self.initialize() async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM apscheduler_jobs WHERE id = ?", (job_id,) ) as cursor: row = await cursor.fetchone() if row: return { "id": row["id"], "next_run_time": row["next_run_time"], "job_state": json.loads(row["job_state"]), } return None async def add_job( self, job_id: str, next_run_time: float, job_state: Dict[str, Any] ) -> bool: """ Add a new job to the database. Args: job_id (str): Unique identifier for the job next_run_time (float): Unix timestamp for next scheduled run job_state (Dict[str, Any]): Job state dictionary Returns: bool: True if successful, False otherwise """ await self.initialize() try: async with aiosqlite.connect(self.db_path) as db: await db.execute( "INSERT INTO apscheduler_jobs (id, next_run_time, job_state) VALUES (?, ?, ?)", (job_id, next_run_time, json.dumps(job_state)), ) await db.commit() return True except aiosqlite.IntegrityError: logger.error(f"Job with ID {job_id} already exists") return False except Exception as e: logger.error(f"Error adding job: {e}") return False async def update_job( self, job_id: str, next_run_time: float, job_state: Dict[str, Any] ) -> bool: """ Update an existing job in the database. Args: job_id (str): ID of the job to update next_run_time (float): New next run time job_state (Dict[str, Any]): Updated job state Returns: bool: True if successful, False if job not found """ await self.initialize() async with aiosqlite.connect(self.db_path) as db: cursor = await db.execute( """UPDATE apscheduler_jobs SET next_run_time = ?, job_state = ? WHERE id = ?""", (next_run_time, json.dumps(job_state), job_id), ) await db.commit() return cursor.rowcount > 0 async def remove_job(self, job_id: str) -> bool: """ Remove a job from the database. Args: job_id (str): ID of the job to remove Returns: bool: True if job was removed, False if not found """ await self.initialize() async with aiosqlite.connect(self.db_path) as db: cursor = await db.execute( "DELETE FROM apscheduler_jobs WHERE id = ?", (job_id,) ) await db.commit() return cursor.rowcount > 0 async def remove_all_jobs(self) -> int: """ Remove all jobs from the database. Returns: int: Number of jobs removed """ await self.initialize() async with aiosqlite.connect(self.db_path) as db: cursor = await db.execute("DELETE FROM apscheduler_jobs") await db.commit() return cursor.rowcount async def get_due_jobs(self, now: float) -> List[Dict[str, Any]]: """ Get all jobs that are due to run. Args: now (float): Current timestamp Returns: List[Dict[str, Any]]: List of due jobs """ await self.initialize() async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( """SELECT * FROM apscheduler_jobs WHERE next_run_time <= ? ORDER BY next_run_time""", (now,), ) as cursor: rows = await cursor.fetchall() return [ { "id": row["id"], "next_run_time": row["next_run_time"], "job_state": json.loads(row["job_state"]), } for row in rows ]