- Logger doesn't start with log level DEBUG by default, instead reads a environment variable - Added extensive debug logging - Wrote a readme - Changed database healthcheck loop to only catch pymsql errors
94 lines
2.8 KiB
Python
94 lines
2.8 KiB
Python
from contextlib import asynccontextmanager
|
|
import aiomysql
|
|
import asyncio
|
|
from secret_handler import return_credentials
|
|
import os
|
|
from logger_handler import setup_logger
|
|
import pymysql.err
|
|
|
|
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
|
|
logger.debug(f"[DB] Initialized DBManager: host={host}, db={db}, port={port}, pool_size={pool_size}")
|
|
|
|
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")
|
|
logger.debug(f"[DB] Pool object: {self._pool}")
|
|
self._health_task = asyncio.create_task(self._healthcheck_loop())
|
|
logger.debug("[DB] Healthcheck task started")
|
|
|
|
async def _healthcheck_loop(self):
|
|
logger.debug("[DB] Healthcheck loop running")
|
|
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] Pool healthcheck succeeded")
|
|
except pymysql.err.Error as e:
|
|
logger.warning(f"[DB] Healthcheck failed: {e}")
|
|
|
|
@asynccontextmanager
|
|
async def acquire(self):
|
|
logger.debug("[DB] Acquiring connection from pool")
|
|
conn = await self._pool.acquire()
|
|
logger.debug(f"[DB] Connection acquired: {conn}")
|
|
try:
|
|
yield conn
|
|
finally:
|
|
self._pool.release(conn)
|
|
logger.debug(f"[DB] Connection released: {conn}")
|
|
|
|
async def release(self, conn):
|
|
if self._pool:
|
|
self._pool.release(conn)
|
|
logger.debug(f"[DB] Connection manually released: {conn}")
|
|
|
|
async def close(self):
|
|
logger.info("[DB] Closing connection pool...")
|
|
self._closing = True
|
|
if self._health_task and not self._health_task.done():
|
|
self._health_task.cancel()
|
|
try:
|
|
await self._health_task
|
|
logger.debug("[DB] Healthcheck task cancelled")
|
|
except asyncio.CancelledError:
|
|
logger.debug("[DB] Healthcheck task cancellation caught")
|
|
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)
|
|
|
|
|