cleveragents-core/tools/controller/db/session.py

"""DB engine + session helpers.

The controller talks to SQLite (tests + local dev) or Postgres
(multi-machine production deploy). Engine selection is via
``CLEVERAGENTS_DB_URL`` env var; no URL → in-memory SQLite (test
default).

Connection lifetime: tests + master use short-lived sessions per
operation via ``session_scope`` (transactional + auto-rollback on
exception). Worker controllers (Phase 1c) hold longer-lived sessions
for the dequeue + heartbeat loop; that pattern lives in worker.py.
"""
from __future__ import annotations

import os
from collections.abc import Iterator
from contextlib import contextmanager

from sqlalchemy import Engine, create_engine
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

from .models import Base


DEFAULT_DB_URL = "sqlite:///:memory:"


def _resolve_db_url(db_url: str | None = None) -> str:
    """Resolve the DB URL with env-var fallback."""
    if db_url:
        return db_url
    return os.environ.get("CLEVERAGENTS_DB_URL", DEFAULT_DB_URL)


def build_engine(db_url: str | None = None) -> Engine:
    """Create a SQLAlchemy engine. Sets per-dialect tuning:

    - SQLite: enable foreign keys, WAL mode, busy_timeout. ``echo``
      defaults to False; flip via ``CLEVERAGENTS_DB_ECHO=1``.
    - Postgres: pool_pre_ping=True so dead connections are noticed
      before a tick fires against a stale conn.
    """
    url = _resolve_db_url(db_url)
    echo = os.environ.get("CLEVERAGENTS_DB_ECHO", "0").lower() in {"1", "true", "yes"}
    if url.startswith("sqlite"):
        # SQLite-specific tuning: same-thread restriction off (the
        # controller's worker pool uses threads) + reasonable busy
        # timeout. ``check_same_thread=False`` is safe because each
        # thread acquires its own Session from sessionmaker; SQLAlchemy
        # serializes per-connection work.
        #
        # For in-memory SQLite (``:memory:``), StaticPool routes every
        # connection to the same underlying DB so the heartbeat thread
        # + the runner's write + the dequeue all see the same data.
        # Without it, ``:memory:`` gives each connection its own
        # DB — the test harness sees ghost rows.
        is_in_memory = ":memory:" in url
        connect_args: dict[str, object] = {
            "check_same_thread": False, "timeout": 30,
        }
        engine_kwargs: dict[str, object] = {"echo": echo, "connect_args": connect_args}
        if is_in_memory:
            engine_kwargs["poolclass"] = StaticPool
        engine = create_engine(url, **engine_kwargs)
        # WAL + foreign keys + reasonable defaults on every new
        # connection. The ``connect`` event fires per-connection.
        from sqlalchemy import event

        @event.listens_for(engine, "connect")
        def _sqlite_pragmas(conn, _record):
            cur = conn.cursor()
            cur.execute("PRAGMA foreign_keys = ON")
            cur.execute("PRAGMA journal_mode = WAL")
            cur.execute("PRAGMA synchronous = NORMAL")
            cur.execute("PRAGMA busy_timeout = 5000")
            cur.close()

        # Note on transaction isolation: SQLite defaults to
        # BEGIN DEFERRED, which means concurrent dequeue queries
        # (``UPDATE … WHERE attempt_id = (SELECT … LIMIT 1)``) can
        # race — two workers may both pick the inner SELECT before
        # either UPDATE commits, and the loser hits SQLITE_BUSY (and
        # waits up to ``busy_timeout``). This is **acceptable for
        # single-host dev** (where SQLite is the only option) but is
        # the reason MULTI-MACHINE deployments MUST use Postgres
        # (which has ``FOR UPDATE SKIP LOCKED`` — see db/dequeue.py
        # for the dual-dialect path). The RUNBOOK calls this out
        # under "Multiple machines, but one machine takes all the
        # work."
        return engine
    # Postgres + others
    return create_engine(url, echo=echo, pool_pre_ping=True)


def create_all(engine: Engine | None = None) -> None:
    """Create all controller tables on ``engine``. Idempotent — safe
    to call on every controller startup; existing tables are not
    re-created.

    For schema-migration work (alembic), use ``alembic upgrade head``
    instead. v1 ships with ``create_all`` since the schema isn't
    versioned yet.
    """
    if engine is None:
        engine = build_engine()
    Base.metadata.create_all(engine)


@contextmanager
def session_scope(engine: Engine | None = None) -> Iterator[Session]:
    """Transactional session context.

    Usage:
        with session_scope(engine) as session:
            session.add(workflow)
            # auto-commits on clean exit; rolls back on exception.

    Each ``with`` block gets a fresh session. SQLAlchemy's session
    factory handles connection pooling.
    """
    if engine is None:
        engine = build_engine()
    factory = sessionmaker(engine, expire_on_commit=False)
    session = factory()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()