Skip to content

Commit

Permalink
fix: add another exception handler to make sure it crashes if there i…
Browse files Browse the repository at this point in the history
…s any connection issue on session creation
  • Loading branch information
robcxyz committed Jul 19, 2024
1 parent 872330f commit 7b8f5a3
Showing 1 changed file with 13 additions and 4 deletions.
17 changes: 13 additions & 4 deletions balanced_backend/db.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from loguru import logger
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlmodel import SQLModel, create_engine
Expand All @@ -16,7 +17,9 @@
ASYNC_SQLALCHEMY_DATABASE_URL = "postgresql+asyncpg" + SQLALCHEMY_DATABASE_URL_STUB
SQLALCHEMY_DATABASE_URL = "postgresql+psycopg2" + SQLALCHEMY_DATABASE_URL_STUB

logger.info(f"Connecting to server: {settings.POSTGRES_SERVER} and {settings.POSTGRES_DATABASE}")
logger.info(
f"Connecting to server: {settings.POSTGRES_SERVER} and {settings.POSTGRES_DATABASE}"
)

async_engine = create_async_engine(
ASYNC_SQLALCHEMY_DATABASE_URL,
Expand All @@ -34,9 +37,15 @@ async def init_db():


async def get_session() -> AsyncSession:
async_session = sessionmaker(async_engine, class_=AsyncSession, expire_on_commit=False)
async with async_session() as session:
yield session
async_session = sessionmaker(
async_engine, class_=AsyncSession, expire_on_commit=False
)
try:
async with async_session() as session:
yield session
except SQLAlchemyError as e:
logger.error(f"Failed to get a database session: {e}")
raise SystemExit(e)


engine = create_engine(SQLALCHEMY_DATABASE_URL)
Expand Down

0 comments on commit 7b8f5a3

Please sign in to comment.