78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
"""SQLite database setup (SQLAlchemy). Tuned for bulk transaction inserts."""
|
|
from __future__ import annotations
|
|
|
|
from sqlalchemy import create_engine, event
|
|
from sqlalchemy.orm import declarative_base, sessionmaker
|
|
|
|
from ..config import DB_PATH, ensure_dirs
|
|
|
|
ensure_dirs()
|
|
|
|
ENGINE = create_engine(
|
|
f"sqlite:///{DB_PATH}",
|
|
connect_args={"check_same_thread": False},
|
|
future=True,
|
|
)
|
|
|
|
|
|
@event.listens_for(ENGINE, "connect")
|
|
def _sqlite_pragmas(dbapi_conn, _rec):
|
|
cur = dbapi_conn.cursor()
|
|
cur.execute("PRAGMA journal_mode=WAL")
|
|
cur.execute("PRAGMA synchronous=NORMAL")
|
|
cur.execute("PRAGMA foreign_keys=ON")
|
|
cur.execute("PRAGMA cache_size=-64000") # ~64 MB page cache
|
|
cur.execute("PRAGMA busy_timeout=30000") # wait (don't error) when another writer holds the lock
|
|
cur.close()
|
|
|
|
|
|
SessionLocal = sessionmaker(bind=ENGINE, autoflush=False, expire_on_commit=False, future=True)
|
|
Base = declarative_base()
|
|
|
|
|
|
def init_db() -> None:
|
|
from . import models # noqa: F401 (register models)
|
|
Base.metadata.create_all(ENGINE)
|
|
_migrate()
|
|
|
|
|
|
def _migrate() -> None:
|
|
"""Add columns introduced after a DB was first created (SQLite create_all won't)."""
|
|
added = {
|
|
"sessions": [
|
|
("progress_rows_done", "INTEGER DEFAULT 0"),
|
|
("progress_rows_total", "INTEGER DEFAULT 0"),
|
|
("eta_seconds", "INTEGER DEFAULT 0"),
|
|
("opening_mode", "VARCHAR DEFAULT 'zero'"),
|
|
("opening_source_session_id", "INTEGER"),
|
|
],
|
|
"reconciliation": [
|
|
("received_payouts", "FLOAT DEFAULT 0"),
|
|
("all_payouts", "FLOAT DEFAULT 0"),
|
|
],
|
|
"finance_control": [
|
|
("tolerance", "FLOAT DEFAULT 1.0"),
|
|
],
|
|
"exports": [
|
|
("kind", "VARCHAR DEFAULT 'full'"),
|
|
],
|
|
"transactions": [
|
|
("txn_type_en", "VARCHAR"),
|
|
("storage_flag", "BOOLEAN DEFAULT 0"),
|
|
],
|
|
}
|
|
with ENGINE.begin() as conn:
|
|
for table, cols in added.items():
|
|
existing = {r[1] for r in conn.exec_driver_sql(f"PRAGMA table_info({table})")}
|
|
for name, decl in cols:
|
|
if name not in existing:
|
|
conn.exec_driver_sql(f"ALTER TABLE {table} ADD COLUMN {name} {decl}")
|
|
|
|
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|