stackpilot/backend/database.py

"""SQLModel database setup."""
from __future__ import annotations

import logging
import os
from collections.abc import Generator
from typing import Optional

from sqlalchemy import inspect, text
from sqlalchemy.exc import OperationalError
from sqlmodel import Session, SQLModel, create_engine

from config import settings

logger = logging.getLogger("stackpilot.database")

os.makedirs(settings.DATA_DIR, exist_ok=True)
_DB_PATH = os.path.join(settings.DATA_DIR, "stackpilot.db")
_DB_URL = f"sqlite:///{_DB_PATH}"

engine = create_engine(
    _DB_URL,
    echo=False,
    connect_args={"check_same_thread": False},
)


def _default_literal(col) -> Optional[str]:
    """SQL literal for a column's scalar default, or None if it has none.

    Only plain values are rendered — a callable default (``default_factory``,
    e.g. a timestamp) has no fixed literal, so those columns are added nullable
    as before and filled by the ORM on the next write.
    """
    default = col.default
    if default is None or not getattr(default, "is_scalar", False):
        return None
    value = default.arg
    if isinstance(value, bool):
        return "1" if value else "0"
    if isinstance(value, (int, float)):
        return str(value)
    if isinstance(value, str):
        escaped = value.replace("'", "''")
        return f"'{escaped}'"
    return None


def _ensure_model_columns() -> None:
    """Add columns that models define but a pre-existing table is missing.

    ``SQLModel.create_all`` creates missing *tables* but never ALTERs an
    existing one, so installs that predate a newly-added column keep the old
    schema — and every ORM query that names the column fails with
    ``OperationalError: no such column``. For each mapped table we diff the
    model's columns against the live table and ``ADD COLUMN`` the safe
    (nullable, or defaulted) ones. Idempotent: on a fresh DB create_all already
    made every column, so this is a no-op.

    Requires ``models`` to have been imported, or ``SQLModel.metadata`` is empty
    and this silently does nothing. :func:`init_db` imports it first.

    A column with a scalar default is added ``NOT NULL DEFAULT <value>`` so
    existing rows are backfilled in the same statement. Without that clause
    SQLite fills them with NULL, which is how a new non-nullable field turns
    into a runtime surprise — for ``User.token_version`` it would have meant
    every existing session failing its version check after the upgrade.
    """
    insp = inspect(engine)
    live_tables = set(insp.get_table_names())
    with engine.begin() as conn:
        for table_name, table in SQLModel.metadata.tables.items():
            if table_name not in live_tables:
                continue
            existing = {c["name"] for c in insp.get_columns(table_name)}
            for col in table.columns:
                if col.name in existing:
                    continue
                # SQLite can only ADD a NOT NULL column if it has a default to
                # backfill existing rows; skip the rest rather than crash.
                if not col.nullable and col.default is None and col.server_default is None:
                    logger.warning(
                        "Cannot auto-add non-nullable column %s.%s (no default); "
                        "manual migration needed", table_name, col.name
                    )
                    continue
                ddl = f'ALTER TABLE "{table_name}" ADD COLUMN "{col.name}" '
                ddl += col.type.compile(dialect=engine.dialect)
                if (literal := _default_literal(col)) is not None:
                    # Backfills existing rows and satisfies SQLite's rule that a
                    # NOT NULL column may only be added together with a default.
                    ddl += f" NOT NULL DEFAULT {literal}"
                conn.execute(text(ddl))
                logger.info("Schema migration: added column %s.%s", table_name, col.name)


#: Tables and columns left behind when the remote-host (agent) integration was
#: removed in 0.48.0. SQLite before 3.35 cannot DROP COLUMN, and the rows are
#: harmless dead weight either way — so the table goes and the columns are only
#: dropped where the SQLite build supports it.
_REMOVED_TABLES = ("agent",)
_REMOVED_COLUMNS = (("autoupdate", "agent_id"), ("backupschedule", "agent_id"))


def _drop_removed_schema() -> None:
    """Clean up schema left over from features that no longer exist.

    Without this an upgraded install keeps an ``agent`` table full of host URLs
    and bearer tokens for a feature that is gone — credentials sitting in the
    database with nothing to use them.

    Each statement runs in its own transaction on purpose: a failed DDL poisons
    the transaction it is in, so sharing one would mean a single unsupported
    DROP COLUMN takes the table drop down with it.
    """
    insp = inspect(engine)
    live = set(insp.get_table_names())

    for table in _REMOVED_TABLES:
        if table not in live:
            continue
        with engine.begin() as conn:
            conn.execute(text(f'DROP TABLE "{table}"'))
        logger.info("Schema migration: dropped obsolete table %s", table)

    for table, column in _REMOVED_COLUMNS:
        if table not in live:
            continue
        if column not in {c["name"] for c in insp.get_columns(table)}:
            continue
        try:
            with engine.begin() as conn:
                conn.execute(text(f'ALTER TABLE "{table}" DROP COLUMN "{column}"'))
            logger.info("Schema migration: dropped obsolete column %s.%s", table, column)
        except OperationalError:
            # SQLite < 3.35 has no DROP COLUMN. The column is nullable and
            # nothing reads it any more, so leaving it is harmless.
            logger.info(
                "Leaving obsolete column %s.%s in place (this SQLite cannot "
                "drop columns); it is unused", table, column
            )


def init_db() -> None:
    # Import models so they are registered on SQLModel.metadata.
    import models  # noqa: F401

    SQLModel.metadata.create_all(engine)
    _ensure_model_columns()
    _drop_removed_schema()


def get_session() -> Generator[Session, None, None]:
    with Session(engine) as session:
        yield session