118 lines
3.9 KiB
Python
118 lines
3.9 KiB
Python
from collections.abc import Generator
|
|
from pathlib import Path
|
|
from shutil import copy2
|
|
from datetime import UTC, datetime
|
|
|
|
from sqlalchemy import Engine, create_engine, event
|
|
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
|
|
|
|
from app.config import get_settings
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
def create_sqlite_engine(database_url: str):
|
|
if database_url.startswith("sqlite:///") and database_url != "sqlite:///:memory:":
|
|
db_path = Path(database_url.removeprefix("sqlite:///"))
|
|
if db_path.parent != Path("."):
|
|
db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
connect_args = {"check_same_thread": False, "timeout": 10}
|
|
engine = create_engine(database_url, connect_args=connect_args)
|
|
engine.dialect.connect_args = connect_args
|
|
|
|
@event.listens_for(engine, "connect")
|
|
def set_sqlite_pragma(dbapi_connection, _connection_record):
|
|
cursor = dbapi_connection.cursor()
|
|
cursor.execute("PRAGMA journal_mode=WAL")
|
|
cursor.close()
|
|
|
|
return engine
|
|
|
|
|
|
def sqlite_file_path(target_engine: Engine = None) -> Path | None:
|
|
target_engine = target_engine or engine
|
|
if target_engine.url.get_backend_name() != "sqlite":
|
|
return None
|
|
database = target_engine.url.database
|
|
if not database or database == ":memory:":
|
|
return None
|
|
return Path(database)
|
|
|
|
|
|
engine = create_sqlite_engine(get_settings().database_url)
|
|
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
|
|
|
|
|
def check_database_integrity(target_engine: Engine = engine) -> str:
|
|
raw_connection = target_engine.raw_connection()
|
|
try:
|
|
cursor = raw_connection.driver_connection.cursor()
|
|
try:
|
|
cursor.execute("PRAGMA integrity_check")
|
|
result = cursor.fetchone()
|
|
finally:
|
|
cursor.close()
|
|
finally:
|
|
raw_connection.close()
|
|
|
|
status = result[0] if result else "missing_result"
|
|
if status != "ok":
|
|
raise RuntimeError(f"SQLite integrity check failed: {status}")
|
|
return status
|
|
|
|
|
|
def checkpoint_sqlite_wal(target_engine: Engine = engine) -> bool:
|
|
if str(target_engine.url) == "sqlite:///:memory:":
|
|
return False
|
|
if target_engine.url.get_backend_name() != "sqlite":
|
|
return False
|
|
|
|
with target_engine.begin() as connection:
|
|
connection.exec_driver_sql("PRAGMA wal_checkpoint(TRUNCATE)")
|
|
return True
|
|
|
|
|
|
def ensure_sqlite_schema_compat(target_engine: Engine = engine) -> None:
|
|
if target_engine.url.get_backend_name() != "sqlite":
|
|
return
|
|
with target_engine.begin() as connection:
|
|
task_columns = {row[1] for row in connection.exec_driver_sql("PRAGMA table_info(tasks)").all()}
|
|
if "current_stage" not in task_columns:
|
|
connection.exec_driver_sql("ALTER TABLE tasks ADD COLUMN current_stage VARCHAR(64)")
|
|
if "last_progress_at" not in task_columns:
|
|
connection.exec_driver_sql("ALTER TABLE tasks ADD COLUMN last_progress_at DATETIME")
|
|
|
|
|
|
def backup_sqlite_files(target_engine: Engine = engine) -> list[Path]:
|
|
db_path = sqlite_file_path(target_engine)
|
|
if db_path is None:
|
|
return []
|
|
backup_dir = db_path.parent / "corrupt-backups"
|
|
backup_dir.mkdir(parents=True, exist_ok=True)
|
|
timestamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S")
|
|
copied: list[Path] = []
|
|
for source in [db_path, db_path.with_name(f"{db_path.name}-wal"), db_path.with_name(f"{db_path.name}-shm")]:
|
|
if not source.exists():
|
|
continue
|
|
target = backup_dir / f"{source.name}.{timestamp}.bak"
|
|
copy2(source, target)
|
|
copied.append(target)
|
|
return copied
|
|
|
|
|
|
def init_db() -> None:
|
|
import app.models # noqa: F401
|
|
|
|
Base.metadata.create_all(engine)
|
|
ensure_sqlite_schema_compat(engine)
|
|
check_database_integrity(engine)
|
|
checkpoint_sqlite_wal(engine)
|
|
|
|
|
|
def get_db_session() -> Generator[Session]:
|
|
with SessionLocal() as session:
|
|
yield session
|