mihir2007/Cyber-Risk

1"""2database.py3-----------4SQLAlchemy 2.0 engine, session factory, and declarative base for the5AI-Powered Continuous Cyber Risk Quantification (CRQ) platform.6 7Configured for production PostgreSQL with resilient connection pooling:8    - pool_size=109    - max_overflow=2010    - pool_pre_ping=True (recovers from stale/dropped database connections)11    12Connection parameters are dynamically read from the `DATABASE_URL` environment13variable, falling back to a local PostgreSQL instance by default.14"""15 16from __future__ import annotations17 18import os19from collections.abc import Generator20from typing import Any21 22from sqlalchemy import create_engine23from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker24 25from urllib.parse import quote_plus26# Dynamic database URL with safe production PostgreSQL default27raw_password = "mihir@1308"  # e.g. "password@1308" or whatever it is28encoded_password = quote_plus(raw_password)29 30DATABASE_URL = f"postgresql+psycopg2://postgres:{encoded_password}@localhost:5432/cyber_risk_db"31# Engine configuration with production connection pooling parameters32engine_kwargs: dict[str, Any] = {33    "echo": False,34    "future": True,35    "pool_pre_ping": True,36}37 38if not DATABASE_URL.startswith("sqlite"):39    engine_kwargs.update(40        {41            "pool_size": 10,42            "max_overflow": 20,43        }44    )45else:46    # Retain safe thread sharing if SQLite is supplied (e.g. in offline unit tests)47    engine_kwargs["connect_args"] = {"check_same_thread": False}48 49engine = create_engine(DATABASE_URL, **engine_kwargs)50 51SessionLocal = sessionmaker(52    autocommit=False,53    autoflush=False,54    bind=engine,55    future=True,56)57 58 59class Base(DeclarativeBase):60    """Declarative base class shared by all ORM models in the platform."""61 62    pass63 64 65def init_db() -> None:66    """Creates all database tables defined in the ORM schema if they do not exist."""67    # Imported locally to avoid circular import between database and models68    import models  # noqa: F40169 70    Base.metadata.create_all(bind=engine)71 72 73def get_db() -> Generator[Session, None, None]:74    """75    FastAPI dependency and context generator that yields a SQLAlchemy session76    and guarantees proper session cleanup and closure.77    """78    db: Session = SessionLocal()79    try:80        yield db81    finally:82        db.close()83