mirror of
https://github.com/Donkie/Spoolman.git
synced 2026-10-22 10:02:17 -05:00
Donkie
72daa29cd2
Task 10: config.js JSON-encodes the base path, so a backslash cannot escape the string literal the way the old " check allowed; search.py escapes LIKE wildcards at both its call sites, using a helper moved into database/utils.py so it and extra_field_query share one implementation. Task 9 was a pair of decisions rather than fixes. /info keeps its paths: both clients render them and removing response fields breaks API v1, while the disclosure buys nothing against an API that is unauthenticated anyway. /metrics keeps its shape but now says what it publishes when switched on. Also pins the unit tests to a throwaway data directory -- importing spoolman.main creates one as a side effect, and it was the real one.
299 lines
11 KiB
Python
299 lines
11 KiB
Python
"""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():
|
|
# /metrics is unauthenticated like the rest of the API, but unlike the rest of it this is
|
|
# opt-in, so the operator is here on purpose and can be told what they just published.
|
|
logger.warning(
|
|
"Metrics are enabled, so /metrics serves your vendor and filament names, colors and "
|
|
"per-spool prices to anyone who can reach it, with no authentication. Keep it off the "
|
|
"public internet, or put it behind your reverse proxy's authentication.",
|
|
)
|
|
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()
|