backend/app/database/chat_migrate.py

"""Lightweight schema updates for chat sessions (SQLite/HF compatible)."""
from sqlalchemy import inspect, text

from app.database.database import engine
from app.database.models import ChatSession, ChatMessage


def migrate_chat_schema() -> None:
    insp = inspect(engine)
    tables = insp.get_table_names()

    if "chat_sessions" not in tables:
        ChatSession.__table__.create(bind=engine, checkfirst=True)

    if "chat_messages" not in tables:
        ChatMessage.__table__.create(bind=engine, checkfirst=True)
        return

    cols = {c["name"] for c in insp.get_columns("chat_messages")}
    if "session_id" not in cols:
        with engine.begin() as conn:
            conn.execute(text("ALTER TABLE chat_messages ADD COLUMN session_id VARCHAR"))

    _attach_orphan_messages()


def _attach_orphan_messages() -> None:
    from sqlalchemy.orm import Session
    from app.database.database import SessionLocal

    db = SessionLocal()
    try:
        orphans = (
            db.query(ChatMessage)
            .filter(ChatMessage.session_id.is_(None))
            .all()
        )
        if not orphans:
            return
        by_user: dict[str, list] = {}
        for msg in orphans:
            by_user.setdefault(msg.user_id, []).append(msg)
        for user_id, msgs in by_user.items():
            session = ChatSession(user_id=user_id, title="Previous conversation")
            db.add(session)
            db.flush()
            for m in msgs:
                m.session_id = session.id
            session.updated_at = msgs[-1].created_at if msgs[-1].created_at else session.updated_at
        db.commit()
    finally:
        db.close()