second-brain/src/second_brain/database.py

"""
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"]