""" Database setup (SQLAlchemy) — MySQL for shared use, SQLite for a laptop or a demo. Which one is used comes from `config.DB_BACKEND`. Everything above this layer is written against SQLAlchemy and is dialect-agnostic; the two places that are not — the raw bulk INSERT in services/store.py and the column-migration below — ask the engine which dialect it is rather than assuming. """ from __future__ import annotations import logging import time from urllib.parse import quote_plus from sqlalchemy import create_engine, event, text from sqlalchemy.orm import declarative_base, sessionmaker from ..config import ( DB_BACKEND, MYSQL_DATABASE, MYSQL_HOST, MYSQL_PASSWORD, MYSQL_POOL_RECYCLE, MYSQL_POOL_SIZE, MYSQL_PORT, MYSQL_SLOW_QUERY_MS, MYSQL_USER, database_label, database_url, ensure_dirs, ) ensure_dirs() logger = logging.getLogger(__name__) IS_MYSQL = DB_BACKEND == "mysql" def _ensure_database() -> None: """Create MYSQL_DATABASE if it does not exist yet (MySQL only).""" user = quote_plus(MYSQL_USER) password = quote_plus(MYSQL_PASSWORD) server_url = ( f"mysql+pymysql://{user}:{password}" f"@{MYSQL_HOST}:{MYSQL_PORT}/?charset=utf8mb4" ) server_engine = create_engine(server_url, isolation_level="AUTOCOMMIT") try: with server_engine.connect() as conn: conn.execute( text( f"CREATE DATABASE IF NOT EXISTS `{MYSQL_DATABASE}` " "CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci" ) ) finally: server_engine.dispose() if IS_MYSQL: _ensure_database() ENGINE = create_engine( database_url(), pool_size=MYSQL_POOL_SIZE, pool_recycle=MYSQL_POOL_RECYCLE, pool_pre_ping=True, future=True, ) else: # check_same_thread=False: processing runs in a background thread with its own session. ENGINE = create_engine( database_url(), future=True, connect_args={"check_same_thread": False, "timeout": 30}, ) @event.listens_for(ENGINE, "connect") def _sqlite_pragmas(dbapi_conn, _rec): cur = dbapi_conn.cursor() cur.execute("PRAGMA journal_mode=WAL") # readers don't block the writer cur.execute("PRAGMA foreign_keys=ON") # cascade deletes behave like MySQL cur.execute("PRAGMA busy_timeout=30000") # bulk insert vs progress updates cur.execute("PRAGMA synchronous=NORMAL") cur.close() logger.info("database: %s", database_label()) if MYSQL_SLOW_QUERY_MS > 0: @event.listens_for(ENGINE, "before_cursor_execute") def _before_cursor_execute(conn, cursor, statement, parameters, context, executemany): conn.info["query_start_time"] = time.perf_counter() @event.listens_for(ENGINE, "after_cursor_execute") def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany): start = conn.info.pop("query_start_time", None) if start is None: return elapsed_ms = (time.perf_counter() - start) * 1000 if elapsed_ms >= MYSQL_SLOW_QUERY_MS: logger.warning( "Slow query (%.0f ms): %s", elapsed_ms, statement[:500], ) 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 (create_all won't alter). MySQL DDL rules that differ from SQLite and silently broke this list during the port: * VARCHAR **must** carry a length — a bare `VARCHAR` is a syntax error. Lengths here must match the model's String(n) or the column ends up a different width. * TEXT/BLOB columns cannot take a literal DEFAULT before MySQL 8.0.13, so `TEXT DEFAULT ''` fails. Declare plain TEXT and let the ORM default apply on insert. """ added = { "users": [ ("reset_code_hash", "VARCHAR(255) DEFAULT ''"), ("reset_code_expires", "DATETIME"), ("reset_code_attempts", "INTEGER DEFAULT 0"), ("is_admin", "BOOLEAN DEFAULT 0"), ], "sessions": [ ("progress_rows_done", "INTEGER DEFAULT 0"), ("progress_rows_total", "INTEGER DEFAULT 0"), ("eta_seconds", "INTEGER DEFAULT 0"), ("opening_mode", "VARCHAR(32) DEFAULT 'zero'"), ("opening_source_session_id", "INTEGER"), ("blocked_reason", "TEXT"), ("payout_mode", "VARCHAR(32) DEFAULT 'auto'"), ("needs_reprocess", "BOOLEAN DEFAULT 0"), ], "session_files": [ ("sheet_last_row", "INTEGER DEFAULT 0"), ("blank_rows_skipped", "INTEGER DEFAULT 0"), ("helper_rows_skipped", "INTEGER DEFAULT 0"), ], "fx_rates": [ ("confirmed_by", "VARCHAR(255) DEFAULT ''"), ("confirmed_at", "DATETIME"), ("confirmed_month", "VARCHAR(32) DEFAULT ''"), ], "journal_entries": [ ("reviewed_by", "VARCHAR(255) DEFAULT ''"), ("reviewed_at", "DATETIME"), ("approved_by", "VARCHAR(255) DEFAULT ''"), ("approved_at", "DATETIME"), ], "reconciliation": [ ("received_payouts", "FLOAT DEFAULT 0"), ("all_payouts", "FLOAT DEFAULT 0"), ], "finance_control": [ ("tolerance", "FLOAT DEFAULT 1.0"), ], "exports": [ ("kind", "VARCHAR(255) DEFAULT 'full'"), ], "transactions": [ ("txn_type_en", "VARCHAR(255)"), ("storage_flag", "BOOLEAN DEFAULT 0"), ], } # Each ALTER runs on its own connection scope: MySQL auto-commits DDL, so wrapping the # whole loop in one transaction gives no rollback anyway — and one bad statement would # otherwise abort every later migration for the rest of the run. with ENGINE.connect() as conn: db_name = conn.execute(text("SELECT DATABASE()")).scalar() if IS_MYSQL else None for table, cols in added.items(): if IS_MYSQL: existing = { r[0] for r in conn.execute( text( "SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS " "WHERE TABLE_SCHEMA = :schema AND TABLE_NAME = :table" ), {"schema": db_name, "table": table}, ) } else: existing = {r[1] for r in conn.execute(text(f"PRAGMA table_info({table})"))} if not existing: continue # table not created yet — create_all owns it for name, decl in cols: if name in existing: continue # SQLite has no VARCHAR length limit and rejects some MySQL type spellings; # its dynamic typing makes the declared type advisory anyway. sql_decl = decl.replace("VARCHAR(255)", "VARCHAR").replace( "VARCHAR(32)", "VARCHAR") if not IS_MYSQL else decl try: conn.execute(text(f"ALTER TABLE `{table}` ADD COLUMN `{name}` {sql_decl}")) conn.commit() except Exception: # noqa: BLE001 conn.rollback() logger.exception("migration failed: %s.%s %s", table, name, decl) def get_db(): db = SessionLocal() try: yield db finally: db.close()