"""Database engine and session management for the Aegis platform. The engine and session factory are created lazily so that tests can override ``DATABASE_URL`` via environment variables before any import triggers real PostgreSQL engine creation (which requires psycopg2). """ # Import Generator from collections.abc from collections.abc import Generator # Import create_engine from sqlalchemy from sqlalchemy import create_engine # Import Engine from sqlalchemy.engine from sqlalchemy.engine import Engine # Import Session, declarative_base, sessionmaker from sqlalchemy.orm from sqlalchemy.orm import Session, declarative_base, sessionmaker # Assign Base = declarative_base() Base = declarative_base() # Engine and session factory are created lazily so that tests can # override DATABASE_URL via environment *before* any import triggers # the real PostgreSQL engine creation (which requires psycopg2). _engine = None # Assign _SessionLocal = None _SessionLocal = None # Define function _get_engine def _get_engine() -> Engine: """Return the shared SQLAlchemy engine, creating it on first call. Returns: Engine: Configured SQLAlchemy engine for the application database. """ # Declare global variable global _engine # Check: _engine is None if _engine is None: # Import settings from app.config from app.config import settings # Assign url = settings.DATABASE_URL url = settings.DATABASE_URL # Assign kwargs = {} kwargs: dict = {} # Check: url.startswith("postgresql") if url.startswith("postgresql"): # Call kwargs.update() kwargs.update( # Keyword argument: pool_size pool_size=20, # Keyword argument: max_overflow max_overflow=10, # Keyword argument: pool_recycle pool_recycle=3600, # Keyword argument: pool_pre_ping pool_pre_ping=True, ) # Assign _engine = create_engine(url, **kwargs) _engine = create_engine(url, **kwargs) # Return _engine return _engine # Define function _get_session_factory def _get_session_factory() -> sessionmaker: """Return the shared sessionmaker, creating it on first call. Returns: sessionmaker: Configured sessionmaker bound to the application engine. """ # Declare global variable global _SessionLocal # Check: _SessionLocal is None if _SessionLocal is None: # Assign _SessionLocal = sessionmaker( _SessionLocal = sessionmaker( # Keyword argument: autocommit autocommit=False, autoflush=False, bind=_get_engine() ) # Return _SessionLocal return _SessionLocal # Define class _LazySessionLocal class _LazySessionLocal: """Proxy so ``SessionLocal()`` keeps working as before but the real sessionmaker is only created on first call.""" # Define function __call__ def __call__(self, *args: object, **kwargs: object) -> Session: """Create and return a new database session. Args: *args (object): Positional arguments forwarded to the sessionmaker. **kwargs (object): Keyword arguments forwarded to the sessionmaker. Returns: Session: A new SQLAlchemy database session. """ # Return _get_session_factory()(*args, **kwargs) return _get_session_factory()(*args, **kwargs) # Define function __getattr__ def __getattr__(self, name: str) -> object: """Delegate attribute access to the underlying sessionmaker. Args: name (str): Attribute name to look up on the sessionmaker. Returns: object: The attribute value from the underlying sessionmaker. """ # Return getattr(_get_session_factory(), name) return getattr(_get_session_factory(), name) # Assign SessionLocal = _LazySessionLocal() SessionLocal = _LazySessionLocal() # Define class _EngineProxy class _EngineProxy: """Thin proxy so ``from app.database import engine`` still works.""" # Define function __getattr__ def __getattr__(self, name: str) -> object: """Delegate attribute access to the lazily-created engine. Args: name (str): Attribute name to look up on the real engine. Returns: object: The attribute value from the underlying SQLAlchemy engine. """ # Return getattr(_get_engine(), name) return getattr(_get_engine(), name) # Assign engine = _EngineProxy() # type: ignore[assignment] engine = _EngineProxy() # type: ignore[assignment] # Define function get_db def get_db() -> Generator[Session, None, None]: """Yield a database session and close it when the request is done. Intended for use as a FastAPI dependency. Yields: Session: An active SQLAlchemy session for the current request. """ # Assign db = SessionLocal() db = SessionLocal() # Attempt the following; catch errors below try: # Yield db yield db # Always execute this cleanup block finally: # Close the database session db.close()