backend/app/core/database.py

from collections.abc import Generator, Iterator
from contextlib import contextmanager
import sqlite3

from sqlalchemy import event
from sqlmodel import Session, SQLModel, create_engine

from app.core.config import get_settings


def _database_url() -> str:
    path = get_settings().database_path
    path.parent.mkdir(parents=True, exist_ok=True)
    return f"sqlite:///{path}"


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


@event.listens_for(engine, "connect")
def _set_sqlite_pragma(dbapi_connection, _connection_record) -> None:
    cursor = dbapi_connection.cursor()
    cursor.execute("PRAGMA foreign_keys=ON")
    cursor.close()


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


@contextmanager
def session_scope() -> Iterator[Session]:
    with Session(engine, expire_on_commit=False) as session:
        yield session


def _migrate_single_user_schema() -> None:
    path = get_settings().database_path
    if not path.exists():
        return
    with sqlite3.connect(path) as connection:
        project_columns = [row[1] for row in connection.execute("PRAGMA table_info(projects)").fetchall()]
        if "user_id" not in project_columns:
            connection.execute("DROP TABLE IF EXISTS users")
            return
        connection.execute("PRAGMA foreign_keys=OFF")
        connection.execute(
            """
            CREATE TABLE projects_new (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                description TEXT NOT NULL DEFAULT '',
                status TEXT NOT NULL DEFAULT 'initialized',
                created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
                last_modified TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
                active_data_file_id INTEGER
            )
            """
        )
        connection.execute(
            """
            INSERT INTO projects_new (
                id,
                name,
                description,
                status,
                created_at,
                last_modified,
                active_data_file_id
            )
            SELECT
                id,
                name,
                description,
                status,
                created_at,
                last_modified,
                active_data_file_id
            FROM projects
            """
        )
        connection.execute("DROP TABLE projects")
        connection.execute("ALTER TABLE projects_new RENAME TO projects")
        connection.execute("DROP TABLE IF EXISTS users")
        connection.execute("PRAGMA foreign_keys=ON")


def init_db() -> None:
    _migrate_single_user_schema()
    SQLModel.metadata.create_all(engine)