Finance-Accounts/ar-aging-app/backend/app/db/database.py

"""
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 = {
        "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()