studspider/src/scheduler.py

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
                ]