backend/src/sona_ai/db/engine.py

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

from sona_ai.core import PROJECT_ROOT


DB_PATH = PROJECT_ROOT / "backend" / "data" / "sona.db"
DB_PATH.parent.mkdir(parents=True, exist_ok=True)

engine = create_engine(
    f"sqlite:///{DB_PATH}",
    connect_args={"check_same_thread": False},
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


class Base(DeclarativeBase):
    pass


@event.listens_for(engine, "connect")
def _enable_sqlite_foreign_keys(dbapi_connection, connection_record):
    cursor = dbapi_connection.cursor()
    cursor.execute("PRAGMA foreign_keys=ON")
    cursor.close()


def init_db():
    from sona_ai.db import models  # noqa: F401

    Base.metadata.create_all(bind=engine)
    _migrate_sqlite()


def _migrate_sqlite():
    with engine.begin() as connection:
        columns = {
            row[1]
            for row in connection.execute(text("PRAGMA table_info(recordings)"))
        }
        summary_columns = {
            row[1]
            for row in connection.execute(text("PRAGMA table_info(recording_summaries)"))
        }
        if "format_name" not in summary_columns:
            connection.execute(
                text("ALTER TABLE recording_summaries ADD COLUMN format_name VARCHAR(255)")
            )
        if "plan_json" not in summary_columns:
            connection.execute(
                text("ALTER TABLE recording_summaries ADD COLUMN plan_json TEXT")
            )
        if "strategy" not in summary_columns:
            connection.execute(
                text("ALTER TABLE recording_summaries ADD COLUMN strategy VARCHAR(32) NOT NULL DEFAULT 'adaptive'")
            )
        if "device" not in columns:
            connection.execute(
                text("ALTER TABLE recordings ADD COLUMN device VARCHAR(32) NOT NULL DEFAULT 'auto'")
            )
        if "processing_stage" not in columns:
            connection.execute(
                text("ALTER TABLE recordings ADD COLUMN processing_stage VARCHAR(64)")
            )
        if "processing_job_id" not in columns:
            connection.execute(
                text("ALTER TABLE recordings ADD COLUMN processing_job_id VARCHAR(36)")
            )
        if "progress_completed_steps" not in columns:
            connection.execute(
                text("ALTER TABLE recordings ADD COLUMN progress_completed_steps INTEGER NOT NULL DEFAULT 0")
            )
        if "progress_total_steps" not in columns:
            connection.execute(
                text("ALTER TABLE recordings ADD COLUMN progress_total_steps INTEGER NOT NULL DEFAULT 0")
            )