84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
import asyncio
|
|
from typing import Any, Dict, List, Optional, Type, TypeVar, Union
|
|
|
|
import aiomysql
|
|
from loguru import logger
|
|
|
|
from data import config
|
|
from ..basestorage.storage import RawConnection
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
# noinspection DuplicatedCode
|
|
class MysqlConnection(RawConnection):
|
|
connection_pool = None
|
|
|
|
@staticmethod
|
|
async def __make_request(
|
|
sql: str,
|
|
params: Union[tuple, List[tuple]] = None,
|
|
fetch: bool = False,
|
|
mult: bool = False,
|
|
retries_count: int = 5
|
|
) -> Optional[Union[List[Dict[str, Any]], Dict[str, Any]]]:
|
|
if MysqlConnection.connection_pool is None:
|
|
MysqlConnection.connection_pool = await aiomysql.create_pool(**config.mysql_info)
|
|
async with MysqlConnection.connection_pool.acquire() as conn:
|
|
conn: aiomysql.Connection = conn
|
|
async with conn.cursor(aiomysql.DictCursor) as cur:
|
|
cur: aiomysql.DictCursor = cur
|
|
for i in range(retries_count):
|
|
try:
|
|
if isinstance(params, list):
|
|
await cur.executemany(sql, params)
|
|
else:
|
|
await cur.execute(sql, params)
|
|
except (aiomysql.OperationalError, aiomysql.InternalError) as e:
|
|
logger.error(f'Found error [{e}] [{sql}] [{params}] retrying [{i}/{retries_count}]')
|
|
if 'Deadlock found' in str(e):
|
|
await asyncio.sleep(1)
|
|
else:
|
|
break
|
|
if fetch:
|
|
if mult:
|
|
r = await cur.fetchall()
|
|
else:
|
|
r = await cur.fetchone()
|
|
return r
|
|
else:
|
|
await conn.commit()
|
|
|
|
@staticmethod
|
|
def _convert_to_model(data: Optional[dict], model: Type[T]) -> Optional[T]:
|
|
if data is not None:
|
|
return model(**data)
|
|
else:
|
|
return None
|
|
|
|
@staticmethod
|
|
async def _make_request(
|
|
sql: str,
|
|
params: Union[tuple, List[tuple]] = None,
|
|
fetch: bool = False,
|
|
mult: bool = False,
|
|
model_type: Type[T] = None
|
|
) -> Optional[Union[List[T], T]]:
|
|
raw = await MysqlConnection.__make_request(sql, params, fetch, mult)
|
|
if raw is None:
|
|
if mult:
|
|
return []
|
|
else:
|
|
return None
|
|
else:
|
|
if mult:
|
|
if model_type is not None:
|
|
return [MysqlConnection._convert_to_model(i, model_type) for i in raw]
|
|
else:
|
|
return list(raw)
|
|
else:
|
|
if model_type is not None:
|
|
return MysqlConnection._convert_to_model(raw, model_type)
|
|
else:
|
|
return raw
|