""" Database engine and session management for second-brain. Postgres + pgvector backend on the petalbrain instance. Resolution order for the connection URL: 1. constructor arg passed to `get_database()` 2. env var `SECOND_BRAIN_DATABASE_URL` 3. env var `HERBYLAB_DATABASE_URL` (shared homelab convention) 4. `[database].url` in settings.toml The URL is expected to use SQLAlchemy's `postgresql+psycopg://` driver (psycopg v3). Plain `postgresql://` URLs are accepted and rewritten to use psycopg; this matches the vault-mcp convention where one URL feeds both Alembic (`+psycopg`) and raw psycopg (no prefix). Connection mode: containerized (homelab-postgres:5432 on the homelab docker network). For local-host dev runs, point at 127.0.0.1:5433 instead. """ from __future__ import annotations import os from contextlib import contextmanager from typing import Generator from sqlalchemy import create_engine, text from sqlalchemy.engine import Engine from sqlalchemy.orm import Session, sessionmaker from second_brain.models import SCHEMA, Base def _normalize_url(url: str) -> str: """Ensure the URL uses SQLAlchemy's psycopg-v3 driver.""" if url.startswith("postgresql+psycopg://"): return url if url.startswith("postgresql://"): return url.replace("postgresql://", "postgresql+psycopg://", 1) return url def _resolve_url(explicit: str | None) -> str: if explicit: return _normalize_url(explicit) for var in ("SECOND_BRAIN_DATABASE_URL", "HERBYLAB_DATABASE_URL"): v = os.environ.get(var) if v: return _normalize_url(v) # Fall back to settings.toml [database].url try: from second_brain.config import load_config cfg = load_config() url = cfg.database.get("url") if url: return _normalize_url(url) except Exception: pass raise RuntimeError( "No Postgres connection URL found. Set SECOND_BRAIN_DATABASE_URL " "(or HERBYLAB_DATABASE_URL), or fill `[database].url` in " "config/settings.toml." ) class Database: """Manages the SQLAlchemy engine and session factory for second-brain.""" def __init__(self, url: str | None = None) -> None: self.url = _resolve_url(url) # Pool sizing matches vault-mcp (min=1/max=10). No PgBouncer; a normal # QueuePool is fine for the CLI + scheduler + web workers we run. self.engine: Engine = create_engine( self.url, echo=False, pool_size=int(os.environ.get("SECOND_BRAIN_DB_POOL_MIN", "1")), max_overflow=int(os.environ.get("SECOND_BRAIN_DB_POOL_MAX", "9")), pool_pre_ping=True, future=True, ) self.SessionLocal = sessionmaker( autocommit=False, autoflush=False, bind=self.engine, future=True, ) def init_db(self) -> None: """Create the schema and tables if missing. Schema creation is idempotent and only relevant for dev/test setups where Alembic hasn't been applied yet. In production, prefer `alembic upgrade head` — this method just defensively ensures the schema exists before `create_all` issues unqualified DDL. """ with self.engine.begin() as conn: conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA}")) Base.metadata.create_all(bind=self.engine) def get_session(self) -> Session: """Return a new bare session (caller must close it).""" return self.SessionLocal() @contextmanager def session(self) -> Generator[Session, None, None]: """Context-manager that commits on exit and rolls back on error.""" sess = self.SessionLocal() try: yield sess sess.commit() except Exception: sess.rollback() raise finally: sess.close() def close(self) -> None: self.engine.dispose() # --------------------------------------------------------------------------- # Singleton # --------------------------------------------------------------------------- _db_instance: Database | None = None def get_database(url: str | None = None) -> Database: """Return the singleton Database, creating it lazily on first call. Schema/tables are *not* auto-created here — run `alembic upgrade head` once during setup. Tests may explicitly call `Database(url).init_db()` against a scratch database. """ global _db_instance if _db_instance is None: _db_instance = Database(url) return _db_instance def reset_database_singleton() -> None: """Test hook — drop the cached singleton so a new URL can take effect.""" global _db_instance if _db_instance is not None: _db_instance.close() _db_instance = None __all__ = ["Database", "get_database", "reset_database_singleton"]