Model/backend/app/db/connection.py

from sqlalchemy import create_engine
from contextlib import contextmanager
from backend.app.config import get_settings
from sqlmodel import Session

connection_string = (
    "postgresql+{drivername}://{username}:{password}@{server}:{port}/{dbname}"
)
db_string = connection_string.format(
    drivername="psycopg2",  # You'll need to use psycopg2 driver for PostgreSQL
    username=get_settings().DB_USERNAME,
    password=get_settings().DB_PASSWORD,
    server=get_settings().DB_HOST,
    port=get_settings().DB_PORT,
    dbname=get_settings().DB_NAME,
)

# db_engine = create_engine(db_string, pool_size=5, max_overflow=5)

# Adjusted database connection to decease pool size for serverless environments (from lambda) so that
# each lambda doesn't hog all connections
db_engine = create_engine(
    db_string,
    pool_size=3,
    max_overflow=5,  # Limit the number of extra connections. With this and pool size, we allow 1 connection per lambda
    pool_pre_ping=True,
    pool_recycle=300,  # Forces SQLAlchemy to close and reopen any connection older than 300 seconds
)


def get_db_session():
    if db_engine is None:
        raise RuntimeError(
            "Database is not configured. Set DATABASE_URL in environment variables."
        )
    return Session(db_engine)


@contextmanager
def db_session():
    session = Session(db_engine)
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager
def db_read_session():
    session = Session(db_engine, expire_on_commit=False)
    try:
        yield session
    finally:
        session.close()