Spoolman/spoolman/database/database.py

"""SQLAlchemy database setup."""

import datetime
import filecmp
import logging
import shutil
import sqlite3
import time
from collections.abc import AsyncGenerator
from os import PathLike
from pathlib import Path
from typing import NamedTuple

from scheduler.asyncio.scheduler import Scheduler
from sqlalchemy import URL
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine

from spoolman import env
from spoolman.prometheus.metrics import filament_metrics, spool_metrics

logger = logging.getLogger(__name__)

# Rotating the backup history is destructive: it deletes the oldest restore point. Enough calls in
# quick succession therefore leave nothing but snapshots of the damage, which is what made a
# bodyless cross-origin form post to /backup worth defending against. Cheap second layer behind
# the origin guard: refuse to rotate more often than this, and hand back the newest backup instead.
MIN_SECONDS_BETWEEN_ROTATIONS = 5 * 60

# Name of the newest backup. Older ones get a .1 ... .N suffix.
BACKUP_NAME = "spoolman.db"


class BackupResult(NamedTuple):
    """The outcome of a backup request."""

    path: Path | None
    """The newest backup, or None if the database cannot be backed up."""

    created: bool
    """Whether a new backup was written. False when an existing one was reused."""


def get_connection_url() -> URL:
    """Construct the connection URL for the database based on environment variables."""
    db_type = env.get_database_type()
    host = env.get_host()
    port = env.get_port()
    database = env.get_database()
    query = env.get_query()
    username = env.get_username()
    password = env.get_password()

    if db_type is None:
        db_type = env.DatabaseType.SQLITE

        database = str(env.get_data_dir().joinpath("spoolman.db"))
        logger.info('No database type specified, using a default SQLite database located at "%s"', database)
    elif db_type is env.DatabaseType.SQLITE:
        if database is not None:
            raise ValueError("Cannot specify a database name when using SQLite.")

        database = str(env.get_data_dir().joinpath("spoolman.db"))
        logger.info('Using SQLite database located at "%s"', database)
    else:
        logger.info('Connecting to database of type "%s" at "%s:%s"', db_type, host, port)

    return URL.create(
        drivername=db_type.to_drivername(),
        host=host,
        port=port,
        database=database,
        query=query or {},
        username=username,
        password=password,
    )


class Database:
    connection_url: URL
    engine: AsyncEngine | None
    session_maker: async_sessionmaker[AsyncSession] | None

    def __init__(self, connection_url: URL) -> None:
        """Construct the Database wrapper and set config parameters."""
        self.connection_url = connection_url
        # Monotonic timestamp of the last rotation, for the rate limit. Process-local, which is
        # fine: Spoolman is single-instance, and a restart erring towards one extra backup is safe.
        self._last_rotation: float | None = None

    def is_file_based_sqlite(self) -> bool:
        """Return True if the database is file based."""
        return (
            self.connection_url.drivername[:6] == "sqlite"
            and self.connection_url.database is not None
            and self.connection_url.database != ":memory:"
        )

    def connect(self) -> None:
        """Connect to the database."""
        if env.get_logging_level() == logging.DEBUG:
            logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)

        connect_args = {}
        if self.connection_url.drivername == "sqlite+aiosqlite":
            connect_args["timeout"] = 60
        connection_options = {}
        if self.connection_url.drivername == "mysql+aiomysql":
            connection_options["pool_recycle"] = 3600
        self.engine = create_async_engine(
            self.connection_url,
            connect_args=connect_args,
            pool_pre_ping=True,
            **connection_options,
        )
        self.session_maker = async_sessionmaker(self.engine, autocommit=False, autoflush=True, expire_on_commit=False)

    def backup(self, target_path: str | PathLike[str]) -> None:
        """Backup the database."""
        if not self.is_file_based_sqlite() or self.connection_url.database is None:
            return

        logger.info("Backing up SQLite database to %s", target_path)

        def progress(_: int, remaining: int, total: int) -> None:
            logger.info("Copied %d of %d pages.", total - remaining, total)

        if self.connection_url.database == target_path:
            raise ValueError("Cannot backup database to itself.")
        if Path(target_path).exists():
            raise ValueError("Backup target file already exists.")

        with sqlite3.connect(self.connection_url.database) as src, sqlite3.connect(target_path) as dst:
            src.backup(dst, pages=1, progress=progress)

        logger.info("Backup complete.")

    def _rotate(self, backup_folder: Path, num_backups: int) -> None:
        """Shift the backup history down by one, discarding the oldest."""
        oldest = backup_folder.joinpath(f"{BACKUP_NAME}.{num_backups}")
        if oldest.exists():
            logger.info("Deleting oldest backup %s", oldest)
            oldest.unlink()

        for i in range(num_backups - 1, -1, -1):
            backup_path = backup_folder.joinpath(BACKUP_NAME if i == 0 else f"{BACKUP_NAME}.{i}")
            if backup_path.exists():
                logger.debug("Rotating backup %s to %s", backup_path, backup_folder.joinpath(f"{BACKUP_NAME}.{i + 1}"))
                shutil.move(backup_path, backup_folder.joinpath(f"{BACKUP_NAME}.{i + 1}"))

    def backup_and_rotate(
        self,
        backup_folder: str | PathLike[str],
        num_backups: int = 5,
    ) -> BackupResult:
        """Backup the database and rotate existing backups.

        Rotation is skipped, and the newest existing backup returned instead, when it happened
        very recently or when the database has not changed since. Both guard the restore history:
        rotation discards the oldest snapshot, so repeated calls would otherwise walk the whole
        history off the end and leave only copies of the current state.

        Args:
            backup_folder: The folder to store the backups in.
            num_backups: The number of backups to keep.

        Returns:
            BackupResult: The newest backup and whether it was created by this call.

        """
        if not self.is_file_based_sqlite() or self.connection_url.database is None:
            logger.info("Skipping backup as the database is not SQLite.")
            return BackupResult(None, created=False)

        backup_folder = Path(backup_folder)
        backup_folder.mkdir(parents=True, exist_ok=True)
        newest = backup_folder.joinpath(BACKUP_NAME)

        if newest.exists() and self._last_rotation is not None:
            since_last = time.monotonic() - self._last_rotation
            if since_last < MIN_SECONDS_BETWEEN_ROTATIONS:
                logger.info(
                    "Skipping backup rotation, the last one was %.0f seconds ago. Returning %s.",
                    since_last,
                    newest,
                )
                return BackupResult(newest, created=False)

        # Write the snapshot first and compare it against the newest backup, rather than comparing
        # the live database file. Under WAL the live file does not yet contain recent commits,
        # so comparing it would skip backups that were genuinely needed.
        pending = backup_folder.joinpath(f"{BACKUP_NAME}.pending")
        pending.unlink(missing_ok=True)
        try:
            self.backup(pending)

            if newest.exists() and filecmp.cmp(pending, newest, shallow=False):
                logger.info("Database is unchanged since %s, keeping the existing backup history.", newest)
                return BackupResult(newest, created=False)

            self._rotate(backup_folder, num_backups)
            shutil.move(pending, newest)
        finally:
            pending.unlink(missing_ok=True)

        self._last_rotation = time.monotonic()
        return BackupResult(newest, created=True)


__db: Database | None = None


def setup_db(connection_url: URL) -> None:
    """Connect to the database.

    Args:
        connection_url: The URL to connect to the database.

    """
    global __db  # noqa: PLW0603
    __db = Database(connection_url)
    __db.connect()


async def backup_global_db(num_backups: int = 5) -> BackupResult:
    """Backup the database and rotate existing backups.

    Returns:
        BackupResult: The newest backup and whether it was created by this call.

    """
    if __db is None:
        raise RuntimeError("DB is not setup.")
    return __db.backup_and_rotate(env.get_backups_dir(), num_backups=num_backups)


async def _backup_task() -> BackupResult:
    """Perform scheduled backup of the database."""
    logger.info("Performing scheduled database backup.")
    if __db is None:
        raise RuntimeError("DB is not setup.")
    return __db.backup_and_rotate(env.get_backups_dir(), num_backups=5)


async def _metrics() -> None:
    """Create some useful prometheus metrics."""
    logger.debug("Start metrics collection")
    async for session in get_db_session():
        await filament_metrics(session)
        await spool_metrics(session)
    logger.debug("End metrics collection")


def schedule_tasks(scheduler: Scheduler) -> None:
    """Schedule tasks to be executed by the provided scheduler.

    Args:
        scheduler: The scheduler to use for scheduling tasks.

    """
    if __db is None:
        raise RuntimeError("DB is not setup.")
    if env.is_metrics_enabled():
        logger.info("Scheduling automatic metric collection.")
        # Run every minute, may be needs specify timer
        scheduler.minutely(datetime.time(second=0), _metrics)  # type: ignore[arg-type]
    if not env.is_automatic_backup_enabled():
        return
    if "sqlite" in __db.connection_url.drivername:
        logger.info("Scheduling automatic database backup for midnight.")
        # Schedule for midnight
        scheduler.daily(datetime.time(hour=0, minute=0, second=0), _backup_task)  # type: ignore[arg-type]


async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
    """Get a DB session to be used with FastAPI's dependency system.

    Yields:
        The database session.

    """
    if __db is None or __db.session_maker is None:
        raise RuntimeError("DB is not setup.")
    async with __db.session_maker() as session:
        try:
            yield session
            await session.commit()
        except Exception as exc:
            await session.rollback()
            raise exc
        finally:
            await session.close()