Spaces:
Running
Running
File size: 2,700 Bytes
89157f5 efba968 89157f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | 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()
|