"""Database utilities for pipeline execution."""

import os
from contextlib import contextmanager
from typing import Any

import psycopg


@contextmanager
def get_database(env_var: str, **kwargs: Any):
    """Get PostgreSQL connection from environment variable.

    Args:
        env_var: Name of environment variable containing the connection string.
        **kwargs: Additional psycopg.connect() arguments.

    Yields:
        psycopg.Connection ready for queries.

    Example:
        with get_database("DATABASE_URL") as conn:
            df = pd.read_sql("SELECT * FROM users", conn)

        with get_database("ANALYTICS_DB") as conn:
            cur = conn.execute("SELECT count(*) FROM events")
            count = cur.fetchone()[0]
    """
    dsn = os.environ.get(env_var)
    if not dsn:
        raise ValueError(f"Environment variable {env_var} not set")

    conn = psycopg.connect(dsn, **kwargs)
    try:
        yield conn
    finally:
        conn.close()
