"""Database connection utilities for Home Assistant PostgreSQL database.""" import os from functools import lru_cache from dotenv import load_dotenv from sqlalchemy import create_engine, text from sqlalchemy.orm import sessionmaker # Load environment variables from .env file load_dotenv() @lru_cache(maxsize=1) def _get_engine(echo: bool = False): """Create and cache the database engine (singleton).""" required_vars = [ "POSTGRES_HOST", "POSTGRES_PORT", "POSTGRES_DB", "POSTGRES_USER", "POSTGRES_PASSWORD", ] missing_vars = [var for var in required_vars if not os.getenv(var)] if missing_vars: raise ValueError(f"Missing required environment variables: {', '.join(missing_vars)}") connection_string = ( f"postgresql://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}" f"@{os.getenv('POSTGRES_HOST')}:{os.getenv('POSTGRES_PORT')}/{os.getenv('POSTGRES_DB')}" ) return create_engine(connection_string, echo=echo) def get_db_session(echo: bool = False) -> sessionmaker: """ Create a SQLAlchemy sessionmaker for the Home Assistant PostgreSQL database. Args: echo: If True, log all SQL statements (default: False) Returns: sessionmaker: Database session maker object Raises: ValueError: If required environment variables are not set """ # Create the SQLAlchemy engine engine = _get_engine(echo=echo) # Create a session maker session = sessionmaker(bind=engine) # Test connection to database with session.begin() as s: s.execute(text("SELECT 1")) return session