Source code for trader.persistence.db

"""Engine, session factory, and schema creation for the PostgreSQL database
(issue #26 — migrated off SQLite; see docs/invariants.md for why)."""

from sqlalchemy import Engine, create_engine

# Re-exported so callers outside this package can catch "the database itself
# refused" — a missing table on an un-migrated database being the common case —
# without importing SQLAlchemy themselves. Narrower than `Exception`, which
# would also swallow the programming errors that must stay loud, and wider
# than `OperationalError`, which misses e.g. `IntegrityError`.
# `__all__` below is what makes this a re-export, so the `as DatabaseError`
# alias the PEP 484 convention would otherwise need is redundant here.
from sqlalchemy.exc import DatabaseError
from sqlalchemy.orm import Session, sessionmaker

from trader.persistence.models import Base

__all__ = [
    "DatabaseError",
    "create_db_engine",
    "create_session_factory",
    "init_db",
]




def create_db_engine(database_url: str) -> Engine:
    """Build an engine for the Postgres database at `database_url`.

    `pool_pre_ping=True`: unlike a SQLite file, Postgres is a real network/
    socket service that can go away mid-process (a `brew services restart`,
    a dropped Unix socket). Without this, a stale pooled connection surfaces
    as a mystifying error mid-cycle instead of being detected and replaced
    before use.

    No PRAGMA setup here (the pre-migration SQLite version of this function
    had one): `journal_mode`/`busy_timeout` have no Postgres equivalent —
    MVCC and its own connection-level lock manager already give concurrent
    readers/writers what WAL mode existed to provide — and `foreign_keys=ON`
    is unconditional on Postgres, never opt-in.
    """
    return create_engine(database_url, pool_pre_ping=True)





def create_session_factory(engine: Engine) -> sessionmaker[Session]:
    """Build a session factory bound to `engine`."""
    return sessionmaker(bind=engine, expire_on_commit=False)





def init_db(database_url: str, stamp: bool = True) -> Engine:
    """Create every table in the (already-existing) Postgres database at
    `database_url`. Safe to run repeatedly.

    Uses `create_all` rather than replaying migrations because the test suite
    builds around a hundred-plus throwaway databases and replaying history for
    each costs far more than it proves. The drift that would otherwise allow
    is caught by `tests/persistence/test_migrations.py`.

    The new database is then stamped at head, so Alembic knows its tables
    already exist and a later `upgrade` applies only genuinely new revisions.
    Pass `stamp=False` to skip that — worth it only for a throwaway database in
    a tight loop, since stamping opens a second connection.

    Unlike the pre-migration SQLite version, this does **not** create the
    database itself — `CREATE DATABASE` is a privileged, out-of-band
    operation (see `tests/persistence/pg_testing.py` for the test-suite
    fixture that does it), never something an app connection triggers
    implicitly by connecting.
    """
    engine = create_db_engine(database_url)
    Base.metadata.create_all(engine)
    if stamp:
        # Imported here, not at module scope: `migrations.py` imports Alembic,
        # and `create_db_engine` is on the hot path for every repository.
        from trader.persistence.migrations import stamp_head

        stamp_head(database_url)
    return engine