82 lines
2.7 KiB
Python
82 lines
2.7 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 _alembic_config(url: str):
|
|
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)
|
|
return cfg
|
|
|
|
|
|
def run_migrations(url: str) -> None:
|
|
"""Upgrade the database at ``url`` to the latest revision."""
|
|
from alembic import command
|
|
|
|
command.upgrade(_alembic_config(url), "head")
|
|
|
|
|
|
def head_revision() -> str | None:
|
|
"""The revision this code expects. ``None`` if the scripts cannot be read."""
|
|
from alembic.script import ScriptDirectory
|
|
|
|
try:
|
|
return ScriptDirectory.from_config(_alembic_config("sqlite://")).get_current_head()
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def current_revision(url: str) -> str | None:
|
|
"""The revision a database is actually at, or ``None`` for an unstamped one."""
|
|
engine = create_db_engine(url)
|
|
try:
|
|
with engine.connect() as connection:
|
|
from alembic.runtime.migration import MigrationContext
|
|
|
|
return MigrationContext.configure(connection).get_current_revision()
|
|
except Exception:
|
|
return None
|
|
finally:
|
|
engine.dispose()
|