62 lines
1.6 KiB
Python
62 lines
1.6 KiB
Python
"""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
|