Spaces:
Sleeping
Sleeping
File size: 2,792 Bytes
57528c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | 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)
|