137 lines
4.0 KiB
Python
137 lines
4.0 KiB
Python
"""MySQL database setup (SQLAlchemy)."""
|
|
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 (
|
|
MYSQL_DATABASE,
|
|
MYSQL_HOST,
|
|
MYSQL_PASSWORD,
|
|
MYSQL_POOL_RECYCLE,
|
|
MYSQL_POOL_SIZE,
|
|
MYSQL_PORT,
|
|
MYSQL_SLOW_QUERY_MS,
|
|
MYSQL_USER,
|
|
ensure_dirs,
|
|
mysql_url,
|
|
)
|
|
|
|
ensure_dirs()
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _ensure_database() -> None:
|
|
"""Create MYSQL_DATABASE if it does not exist yet."""
|
|
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()
|
|
|
|
|
|
_ensure_database()
|
|
|
|
ENGINE = create_engine(
|
|
mysql_url(),
|
|
pool_size=MYSQL_POOL_SIZE,
|
|
pool_recycle=MYSQL_POOL_RECYCLE,
|
|
pool_pre_ping=True,
|
|
future=True,
|
|
)
|
|
|
|
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)."""
|
|
added = {
|
|
"sessions": [
|
|
("progress_rows_done", "INTEGER DEFAULT 0"),
|
|
("progress_rows_total", "INTEGER DEFAULT 0"),
|
|
("eta_seconds", "INTEGER DEFAULT 0"),
|
|
("opening_mode", "VARCHAR(255) 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(255) DEFAULT 'full'"),
|
|
],
|
|
"transactions": [
|
|
("txn_type_en", "VARCHAR(255)"),
|
|
("storage_flag", "BOOLEAN DEFAULT 0"),
|
|
],
|
|
}
|
|
with ENGINE.begin() as conn:
|
|
db_name = conn.execute(text("SELECT DATABASE()")).scalar()
|
|
for table, cols in added.items():
|
|
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},
|
|
)
|
|
}
|
|
for name, decl in cols:
|
|
if name not in existing:
|
|
conn.execute(text(f"ALTER TABLE `{table}` ADD COLUMN `{name}` {decl}"))
|
|
|
|
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|