93 lines
2.6 KiB
Python
93 lines
2.6 KiB
Python
import aiomysql
|
|
import asyncio
|
|
from secret_handler import return_credentials
|
|
import os
|
|
from logger_handler import setup_logger
|
|
|
|
|
|
db_username = return_credentials("/etc/secrets/db_username")
|
|
db_password = return_credentials("/etc/secrets/db_password")
|
|
db_host = os.getenv("BACKEND_PN_DB_HOST","localhost")
|
|
db_database = os.getenv("BACKEND_PN_DB_DATABASE","app")
|
|
|
|
logger = setup_logger(__name__)
|
|
|
|
|
|
class DBManager:
|
|
def __init__(self, host, user, password, db, port=3306, pool_size=5, health_interval=60):
|
|
self.host = host
|
|
self.user = user
|
|
self.password = password
|
|
self.db = db
|
|
self.port = port
|
|
self.pool_size = pool_size
|
|
self._pool: aiomysql.Pool | None = None
|
|
self._health_interval = health_interval
|
|
self._health_task: asyncio.Task | None = None
|
|
self._closing = False
|
|
|
|
async def connect(self):
|
|
self._pool = await aiomysql.create_pool(
|
|
host=self.host,
|
|
user=self.user,
|
|
password=self.password,
|
|
db=self.db,
|
|
port=self.port,
|
|
minsize=1,
|
|
maxsize=self.pool_size,
|
|
autocommit=True,
|
|
cursorclass=aiomysql.DictCursor
|
|
)
|
|
logger.info("[DB] Connection pool created")
|
|
self._health_task = asyncio.create_task(self._healthcheck_loop())
|
|
|
|
async def _healthcheck_loop(self):
|
|
while not self._closing:
|
|
await asyncio.sleep(self._health_interval)
|
|
try:
|
|
async with self.acquire() as conn:
|
|
async with conn.cursor() as cur:
|
|
await cur.execute("SELECT 1")
|
|
logger.debug("[DB] Healthcheck OK")
|
|
except Exception as e:
|
|
logger.warning(f"[DB] Healthcheck failed: {e}")
|
|
|
|
async def acquire(self):
|
|
if not self._pool:
|
|
raise RuntimeError("DB pool not initialized")
|
|
return await self._pool.acquire()
|
|
|
|
async def release(self, conn):
|
|
if self._pool:
|
|
self._pool.release(conn)
|
|
|
|
async def execute(self, query, *args, retries=3):
|
|
for attempt in range(1, retries + 1):
|
|
conn = await self.acquire()
|
|
try:
|
|
async with conn.cursor() as cur:
|
|
await cur.execute(query, args)
|
|
if cur.description:
|
|
return await cur.fetchall()
|
|
return None
|
|
except aiomysql.OperationalError as e:
|
|
logger.warning(f"[DB] Query failed (attempt {attempt}/{retries}): {e}")
|
|
await asyncio.sleep(2 ** (attempt - 1))
|
|
finally:
|
|
await self.release(conn)
|
|
raise RuntimeError("DB query failed after retries")
|
|
|
|
async def close(self):
|
|
self._closing = True
|
|
if self._health_task and not self._health_task.done():
|
|
self._health_task.cancel()
|
|
try:
|
|
await self._health_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
if self._pool:
|
|
self._pool.close()
|
|
await self._pool.wait_closed()
|
|
logger.info("[DB] Connection pool closed")
|
|
|
|
db_manager = DBManager(host=db_host, port=3306, user=db_username, password=db_password, db=db_database) |