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)