From 37d171d198773d4de346b8fb2acc57ee4cb96ee5 Mon Sep 17 00:00:00 2001
From: repair-agent <agent@airlabs.art>
Date: Sat, 28 Feb 2026 16:14:05 +0800
Subject: [PATCH] fix: auto repair bugs #56

---
 app/database.py | 31 +++++++++++++++++++++++++------
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/app/database.py b/app/database.py
index cfcbe2a..cd8d8bd 100644
--- a/app/database.py
+++ b/app/database.py
@@ -3,9 +3,13 @@ from sqlmodel.ext.asyncio.session import AsyncSession
 from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy import text
+from sqlalchemy.exc import DBAPIError, OperationalError
 import os
+import logging
 from dotenv import load_dotenv
 
+logger = logging.getLogger(__name__)
+
 load_dotenv()
 
 DB_USER = os.getenv("DB_USER")
@@ -21,7 +25,10 @@ engine = create_async_engine(
     echo=True,
     future=True,
     pool_pre_ping=True,
-    pool_recycle=300,
+    pool_recycle=180,
+    pool_size=5,
+    max_overflow=10,
+    pool_timeout=30,
 )
 
 async def init_db():
@@ -55,9 +62,21 @@ async def init_db():
             except Exception:
                 pass  # Already applied
 
+_async_session_factory = sessionmaker(
+    engine, class_=AsyncSession, expire_on_commit=False
+)
+
+_MAX_RETRIES = 2
+
 async def get_session() -> AsyncSession:
-    async_session = sessionmaker(
-        engine, class_=AsyncSession, expire_on_commit=False
-    )
-    async with async_session() as session:
-        yield session
+    for attempt in range(_MAX_RETRIES):
+        try:
+            async with _async_session_factory() as 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