53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
"""Database engine, sessions, and migration entry point.
|
|
|
|
SQLite is opened in WAL mode with foreign keys enforced on every connection
|
|
(foreign keys and busy_timeout are per-connection; WAL is persistent but cheap to
|
|
re-assert). Migrations run through Alembic so the same versioned schema is used at
|
|
startup and in tests. Repository and service code depends only on ``Base`` and the
|
|
session factory, never on SQLite specifics, so PostgreSQL stays a future option.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import Engine, create_engine, event
|
|
from sqlalchemy.orm import DeclarativeBase, sessionmaker
|
|
|
|
_REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
def create_db_engine(url: str, *, echo: bool = False) -> Engine:
|
|
connect_args = {"check_same_thread": False} if url.startswith("sqlite") else {}
|
|
engine = create_engine(url, echo=echo, future=True, connect_args=connect_args)
|
|
if url.startswith("sqlite"):
|
|
event.listen(engine, "connect", _apply_sqlite_pragmas)
|
|
return engine
|
|
|
|
|
|
def _apply_sqlite_pragmas(dbapi_connection, _connection_record) -> None:
|
|
cursor = dbapi_connection.cursor()
|
|
cursor.execute("PRAGMA journal_mode=WAL")
|
|
cursor.execute("PRAGMA foreign_keys=ON")
|
|
cursor.execute("PRAGMA busy_timeout=5000")
|
|
cursor.close()
|
|
|
|
|
|
def create_session_factory(engine: Engine) -> sessionmaker:
|
|
return sessionmaker(bind=engine, expire_on_commit=False, future=True)
|
|
|
|
|
|
def run_migrations(url: str) -> None:
|
|
"""Upgrade the database at ``url`` to the latest revision."""
|
|
from alembic import command
|
|
from alembic.config import Config as AlembicConfig
|
|
|
|
cfg = AlembicConfig(str(_REPO_ROOT / "alembic.ini"))
|
|
cfg.set_main_option("script_location", str(_REPO_ROOT / "migrations"))
|
|
cfg.set_main_option("sqlalchemy.url", url)
|
|
command.upgrade(cfg, "head")
|