US01-02: Establish Application and Database Foundation (#47)

photo_pipeline/db.py

@@ -0,0 +1,52 @@
+"""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")