Aegis/backend/app/database.py

"""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()