fix: auto repair bugs #56 #4

Open
zyc wants to merge 1 commits from fix/auto-20260228-161206 into main

View File

@ -3,9 +3,13 @@ from sqlmodel.ext.asyncio.session import AsyncSession
from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.orm import sessionmaker from sqlalchemy.orm import sessionmaker
from sqlalchemy import text from sqlalchemy import text
from sqlalchemy.exc import DBAPIError, OperationalError
import os import os
import logging
from dotenv import load_dotenv from dotenv import load_dotenv
logger = logging.getLogger(__name__)
load_dotenv() load_dotenv()
DB_USER = os.getenv("DB_USER") DB_USER = os.getenv("DB_USER")
@ -21,7 +25,10 @@ engine = create_async_engine(
echo=True, echo=True,
future=True, future=True,
pool_pre_ping=True, pool_pre_ping=True,
pool_recycle=300, pool_recycle=180,
pool_size=5,
max_overflow=10,
pool_timeout=30,
) )
async def init_db(): async def init_db():
@ -55,9 +62,21 @@ async def init_db():
except Exception: except Exception:
pass # Already applied pass # Already applied
async def get_session() -> AsyncSession: _async_session_factory = sessionmaker(
async_session = sessionmaker(
engine, class_=AsyncSession, expire_on_commit=False engine, class_=AsyncSession, expire_on_commit=False
) )
async with async_session() as session:
_MAX_RETRIES = 2
async def get_session() -> AsyncSession:
for attempt in range(_MAX_RETRIES):
try:
async with _async_session_factory() as session:
yield session yield session
return
except (DBAPIError, OperationalError, ConnectionResetError, OSError) as e:
if attempt < _MAX_RETRIES - 1:
logger.warning("Database connection error (attempt %d/%d): %s", attempt + 1, _MAX_RETRIES, e)
await engine.dispose()
continue
raise