import os import boto3 from requests import Session from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from app.models import Base from .constants import AWS_REGION def get_db_credentials(): """Fetch and cache DB credentials from environment or SSM Parameter Store""" if os.getenv("MOCK_AUTH", "").lower() == "true": return ( f"postgresql://{os.getenv('DB_USER')}:{os.getenv('DB_PASSWORD')}" f"@{os.getenv('DB_HOST')}/{os.getenv('DB_NAME')}" ) ssm = boto3.client("ssm", region_name=AWS_REGION) try: host = ssm.get_parameter(Name="/iptv-manager/DB_HOST", WithDecryption=True)[ "Parameter" ]["Value"] user = ssm.get_parameter(Name="/iptv-manager/DB_USER", WithDecryption=True)[ "Parameter" ]["Value"] password = ssm.get_parameter( Name="/iptv-manager/DB_PASSWORD", WithDecryption=True )["Parameter"]["Value"] dbname = ssm.get_parameter(Name="/iptv-manager/DB_NAME", WithDecryption=True)[ "Parameter" ]["Value"] return f"postgresql://{user}:{password}@{host}/{dbname}" except Exception as e: raise RuntimeError(f"Failed to fetch DB credentials from SSM: {str(e)}") # Initialize engine and session maker engine = create_engine(get_db_credentials()) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) def init_db(): """Initialize database by creating all tables""" Base.metadata.create_all(bind=engine) def get_db(): """Dependency for getting database session""" db = SessionLocal() try: yield db finally: db.close() def get_db_session() -> Session: """Get a direct database session (non-generator version)""" return SessionLocal()