fix: auto repair bugs #56

This commit is contained in:
repair-agent 2026-02-28 16:14:05 +08:00
parent 708021c443
commit 37d171d198

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_session_factory = sessionmaker(
engine, class_=AsyncSession, expire_on_commit=False
)
_MAX_RETRIES = 2
async def get_session() -> AsyncSession: async def get_session() -> AsyncSession:
async_session = sessionmaker( for attempt in range(_MAX_RETRIES):
engine, class_=AsyncSession, expire_on_commit=False try:
) async with _async_session_factory() as session:
async with async_session() 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