535fa40ad2
Primarily FastAPI and Pydantic to v2. Also ruff to latest. Updated some code to support Pydantic v2
238 lines
7.9 KiB
Python
238 lines
7.9 KiB
Python
"""SQLAlchemy database setup."""
|
|
|
|
import datetime
|
|
import logging
|
|
import shutil
|
|
import sqlite3
|
|
from collections.abc import AsyncGenerator
|
|
from os import PathLike
|
|
from pathlib import Path
|
|
from typing import Optional, Union
|
|
|
|
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__)
|
|
|
|
|
|
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: Optional[AsyncEngine]
|
|
session_maker: Optional[async_sessionmaker[AsyncSession]]
|
|
|
|
def __init__(self: "Database", connection_url: URL) -> None:
|
|
"""Construct the Database wrapper and set config parameters."""
|
|
self.connection_url = connection_url
|
|
|
|
def is_file_based_sqlite(self: "Database") -> 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: "Database") -> 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
|
|
|
|
self.engine = create_async_engine(
|
|
self.connection_url,
|
|
connect_args=connect_args,
|
|
pool_pre_ping=True,
|
|
)
|
|
self.session_maker = async_sessionmaker(self.engine, autocommit=False, autoflush=True, expire_on_commit=False)
|
|
|
|
def backup(self: "Database", target_path: Union[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 backup_and_rotate(
|
|
self: "Database",
|
|
backup_folder: Union[str, PathLike[str]],
|
|
num_backups: int = 5,
|
|
) -> Optional[Path]:
|
|
"""Backup the database and rotate existing backups.
|
|
|
|
Args:
|
|
backup_folder: The folder to store the backups in.
|
|
num_backups: The number of backups to keep.
|
|
|
|
Returns:
|
|
The path to the created backup or None if no backup was created.
|
|
|
|
"""
|
|
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 None
|
|
|
|
backup_folder = Path(backup_folder)
|
|
backup_folder.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Delete oldest backup
|
|
backup_path = backup_folder.joinpath(f"spoolman.db.{num_backups}")
|
|
if backup_path.exists():
|
|
logger.info("Deleting oldest backup %s", backup_path)
|
|
backup_path.unlink()
|
|
|
|
# Rotate existing backups
|
|
for i in range(num_backups - 1, -1, -1):
|
|
if i == 0:
|
|
backup_path = backup_folder.joinpath("spoolman.db")
|
|
else:
|
|
backup_path = backup_folder.joinpath(f"spoolman.db.{i}")
|
|
if backup_path.exists():
|
|
logger.debug("Rotating backup %s to %s", backup_path, backup_folder.joinpath(f"spoolman.db.{i + 1}"))
|
|
shutil.move(backup_path, backup_folder.joinpath(f"spoolman.db.{i + 1}"))
|
|
|
|
# Create new backup
|
|
backup_path = backup_folder.joinpath("spoolman.db")
|
|
self.backup(backup_path)
|
|
|
|
return backup_path
|
|
|
|
|
|
__db: Optional[Database] = 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) -> Optional[Path]:
|
|
"""Backup the database and rotate existing backups.
|
|
|
|
Returns:
|
|
The path to the created backup or None if no backup was created.
|
|
|
|
"""
|
|
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() -> Optional[Path]:
|
|
"""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()
|