Heimgeist/backend/database.py

from sqlalchemy import text
from sqlalchemy import create_engine
from sqlalchemy.orm import DeclarativeBase, sessionmaker
import uuid

from .paths import database_path, sqlite_database_url

"""
Database utilities and configuration. This module defines the SQLAlchemy
engine, session factory and base class for models. It also contains a
lightweight migration helper used to evolve the schema over time. The
`ensure_sources_column` helper adds the JSON-backed columns used by chat
messages when they do not already exist.

The migration uses SQLite's `ALTER TABLE` syntax and therefore should
only run once on startup. It is safe to call repeatedly: when a column
already exists, the function will simply no-op.
"""

DATABASE_PATH = database_path()
DATABASE_URL = sqlite_database_url(DATABASE_PATH)

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

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

class Base(DeclarativeBase):
    pass


def ensure_sources_column(engine):
    try:
        with engine.begin() as conn:
            cols = [row[1] for row in conn.execute(text("PRAGMA table_info(chat_messages)"))]
            if "message_id" not in cols:
                conn.execute(text("ALTER TABLE chat_messages ADD COLUMN message_id TEXT"))
            if "sources_json" not in cols:
                conn.execute(text("ALTER TABLE chat_messages ADD COLUMN sources_json TEXT DEFAULT '[]'"))
            if "attachments_json" not in cols:
                conn.execute(text("ALTER TABLE chat_messages ADD COLUMN attachments_json TEXT DEFAULT '[]'"))
            for column_name, definition in (
                ("workflow_id", "TEXT"),
                ("workflow_revision_id", "TEXT"),
                ("workflow_run_id", "TEXT"),
                ("agent_summary_json", "TEXT DEFAULT '{}'"),
                ("usage_json", "TEXT DEFAULT '{}'"),
            ):
                if column_name not in cols:
                    conn.execute(text(f"ALTER TABLE chat_messages ADD COLUMN {column_name} {definition}"))
            missing_ids = conn.execute(
                text("SELECT id FROM chat_messages WHERE message_id IS NULL OR message_id = ''")
            ).fetchall()
            for row in missing_ids:
                conn.execute(
                    text("UPDATE chat_messages SET message_id = :message_id WHERE id = :id"),
                    {"message_id": str(uuid.uuid4()), "id": row[0]},
                )
            conn.execute(text(
                "CREATE UNIQUE INDEX IF NOT EXISTS ix_chat_messages_message_id "
                "ON chat_messages (message_id)"
            ))
            conn.execute(text(
                "CREATE INDEX IF NOT EXISTS ix_chat_messages_workflow_run_id "
                "ON chat_messages (workflow_run_id)"
            ))
    except Exception as e:
        print("[db] ensure_sources_column error:", e)