""" Database connection and initialization """ from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker, Session from typing import Generator import logging from app.config import get_settings from app.models.data_models import Base logger = logging.getLogger(__name__) settings = get_settings() # Create database engine engine = create_engine( settings.database_url, echo=settings.debug, pool_pre_ping=True, pool_size=10, max_overflow=20 ) # Create session factory SessionLocal = sessionmaker( autocommit=False, autoflush=False, bind=engine ) def get_db() -> Generator[Session, None, None]: """ Get database session Yields: Database session """ db = SessionLocal() try: yield db finally: db.close() def init_db() -> None: """ Initialize database - create all tables """ try: logger.info("Creating database tables...") Base.metadata.create_all(bind=engine) logger.info("Database tables created successfully!") except Exception as e: logger.error(f"Error creating database tables: {e}") raise def drop_db() -> None: """ Drop all database tables (use with caution!) """ logger.warning("Dropping all database tables...") Base.metadata.drop_all(bind=engine) logger.info("Database tables dropped!") if __name__ == "__main__": # Initialize database when run directly logging.basicConfig(level=logging.INFO) init_db()