241 lines
8.1 KiB
Python
241 lines
8.1 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, connection_url: URL) -> None:
|
|
"""Construct the Database wrapper and set config parameters."""
|
|
self.connection_url = connection_url
|
|
|
|
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: 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,
|
|
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()
|