src/repository/common.py

from typing import Generator
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from contextlib import contextmanager

from dotenv import load_dotenv
import os

load_dotenv()

_engine = None
_SessionLocal = None

def get_engine():
    global _engine
    
    if _engine is None:
        # Load the database URL from environment variables
        database_url = os.getenv("DATABASE_URL")
        if not database_url:
            raise ValueError("DATABASE_URL is not set in environment variables")
        
        # Create the SQLAlchemy engine
        _engine = create_engine(database_url, pool_pre_ping=True)
    
    return _engine

def get_session_local() -> sessionmaker:
    global _SessionLocal

    if _SessionLocal is None:
        # Create a new session local
        _SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=get_engine())
    
    return _SessionLocal

@contextmanager
def get_session() -> Generator[Session, None, None]:
    """
    Get a connection to the Postgres database
    """
    SessionLocal = get_session_local()
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()