database.py

"""
Phase G - database.py
Auto-switches between MySQL (local Docker) and SQLite (HF Spaces).
Set DB_BACKEND=sqlite for HF Spaces, DB_BACKEND=mysql for local Docker.
"""
from __future__ import annotations

import logging
import os
from pathlib import Path

from dotenv import load_dotenv
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase
from typing import AsyncGenerator

load_dotenv(dotenv_path=Path(__file__).parent / ".env", override=True)

log = logging.getLogger("phase_g.database")

DB_BACKEND = os.getenv("DB_BACKEND", "mysql")

if DB_BACKEND == "sqlite":
    SQLITE_PATH  = os.getenv("SQLITE_PATH", "/app/data/nlp.db")
    os.makedirs(os.path.dirname(SQLITE_PATH), exist_ok=True)
    DATABASE_URL = f"sqlite+aiosqlite:///{SQLITE_PATH}"
    engine = create_async_engine(
        DATABASE_URL,
        connect_args={"check_same_thread": False},
        echo=os.getenv("DB_ECHO", "false").lower() == "true",
    )
else:
    from urllib.parse import quote_plus
    MYSQL_USER     = os.getenv("MYSQL_USER", "root")
    MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD", "")
    MYSQL_HOST     = os.getenv("MYSQL_HOST", "localhost")
    MYSQL_PORT     = os.getenv("MYSQL_PORT", "3306")
    MYSQL_DB       = os.getenv("MYSQL_DB", os.getenv("MYSQL_DATABASE", "nlp_classifier"))
    DATABASE_URL = (
        f"mysql+aiomysql://{MYSQL_USER}:{quote_plus(MYSQL_PASSWORD)}"
        f"@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DB}"
        f"?charset=utf8mb4"
    )
    engine = create_async_engine(
        DATABASE_URL,
        pool_size=10,
        max_overflow=20,
        pool_pre_ping=True,
        pool_recycle=1800,
        echo=os.getenv("DB_ECHO", "false").lower() == "true",
    )

AsyncSessionLocal = async_sessionmaker(
    bind=engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autoflush=False,
    autocommit=False,
)


class Base(DeclarativeBase):
    pass


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


async def create_tables() -> None:
    async with engine.begin() as conn:
        await conn.run_sync(lambda conn: Base.metadata.create_all(conn, checkfirst=True))
    log.info("Database tables ready")


async def close_db() -> None:
    await engine.dispose()
    log.info("Database pool closed.")

# rebuild 1781059258.3590245