fix: add some extra try except around db conns #42

icon_governance/db.py

@@ -1,7 +1,10 @@
+import sys
+
 from loguru import logger
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 from sqlalchemy.orm import sessionmaker
 from sqlmodel import SQLModel, create_engine
+from psycopg2 import OperationalError
 
 from icon_governance.config import settings
 from icon_governance.models.preps import Prep
@@ -35,5 +38,9 @@ async def get_session() -> AsyncSession:
         yield session
 
 
-engine = create_engine(SQLALCHEMY_DATABASE_URL)
-session_factory = sessionmaker(bind=engine)
+try:
+    engine = create_engine(SQLALCHEMY_DATABASE_URL)
+    session_factory = sessionmaker(bind=engine)
+except Exception as e:
+    logger.error(f"Database connection dropped: {str(e)}")
+    raise

icon_governance/main_cron.py

@@ -1,3 +1,4 @@
+import sys
 from typing import Callable, TypedDict
 
 from apscheduler.schedulers.background import BlockingScheduler
@@ -105,8 +106,12 @@ class Cron(TypedDict):
 
 
 def run_cron_with_session(cron: Callable):
-    with session_factory() as session:
-        cron(session=session)
+    try:
+        with session_factory() as session:
+            cron(session=session)
+    except Exception as e:
+        logger.error(f"Database operation failed: {str(e)}")
+        sys.exit(1)
 
 
 def main():