menzeljandClaude Opus 4.8
e651029ab2
create_all never ALTERs an existing table, so installs predating the backupschedule.agent_id column kept the old schema and any ORM query naming it failed with "no such column" — which the new fleet dashboard (and the schedules list / scheduler loop) hit. init_db now diffs each mapped table against the live schema and ADD COLUMNs the missing nullable/defaulted ones. Idempotent and self-healing for similar drift. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
"""SQLModel database setup."""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from collections.abc import Generator
|
|
|
|
from sqlalchemy import inspect, text
|
|
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 _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.
|
|
"""
|
|
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_type = col.type.compile(dialect=engine.dialect)
|
|
conn.execute(text(f'ALTER TABLE "{table_name}" ADD COLUMN "{col.name}" {ddl_type}'))
|
|
logger.info("Schema migration: added column %s.%s", table_name, col.name)
|
|
|
|
|
|
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()
|
|
|
|
|
|
def get_session() -> Generator[Session, None, None]:
|
|
with Session(engine) as session:
|
|
yield session
|