app/core/database/mysql.py

from __future__ import annotations

from typing import Any

import aiomysql

from app.core.database.base import BaseExecutor, StatementResult
from app.core.logger import get_logger

_logger = get_logger(__name__)


class MySQLExecutor(BaseExecutor):
    async def _create_pool(self) -> aiomysql.Pool:
        return await aiomysql.create_pool(
            host=self._config.host,
            port=self._config.port,
            user=self._config.username,
            password=self._config.password,
            db=self._config.database,
            minsize=1,
            maxsize=10,
            autocommit=False,
            connect_timeout=self._config.connection_timeout_seconds,
            pool_recycle=3600,
        )

    async def _execute_queries(
        self, pool: aiomysql.Pool, queries: list[str], use_transaction: bool
    ) -> list[StatementResult]:
        async with pool.acquire() as conn:
            if use_transaction:
                await conn.begin()
            try:
                results: list[StatementResult] = []
                for query in queries:
                    result = await self._execute_one(conn, query)
                    results.append(result)
                    if not result.success and use_transaction:
                        await conn.rollback()
                        return results
                if use_transaction:
                    await conn.commit()
                return results
            except Exception:
                if use_transaction:
                    try:
                        await conn.rollback()
                    except Exception:
                        pass
                raise

    async def _execute_one(self, conn: Any, query: str) -> StatementResult:
        try:
            async with conn.cursor(aiomysql.DictCursor) as cursor:
                await cursor.execute(query)
                if cursor.description:
                    rows = await cursor.fetchall()
                    data = [dict(row) for row in rows][: self._config.max_rows]
                    return StatementResult(success=True, rows=len(data), data=data)
                await conn.commit()
                return StatementResult(success=True, rows=cursor.rowcount, data=[])
        except Exception as exc:
            error_code = getattr(exc, "args", [None])[0]
            if isinstance(error_code, int):
                error_code = str(error_code)
            else:
                error_code = type(exc).__name__
            return StatementResult(success=False, error=str(exc), error_code=error_code)

    async def _close_pool(self, pool: aiomysql.Pool) -> None:
        pool.close()
        await pool.wait_closed()