"""RLS (Row-Level Security) helper for Airflow plugin endpoints.

Provides a context manager that opens a sync SQLAlchemy connection,
sets the RLS tenant context via SET LOCAL, and handles commit/rollback.

NOTE: This module runs inside the Airflow api-server process.
Do NOT import from src.*.
"""

from contextlib import contextmanager

from sqlalchemy import text

from api.db import get_engine


@contextmanager
def rls_connection(tenant_id: str):
    """Open a connection with RLS tenant context set.

    Yields the connection inside a transaction. Commits on success,
    rolls back on exception. Always closes the connection on exit.
    """
    engine = get_engine()
    conn = engine.connect()
    txn = conn.begin()
    try:
        conn.execute(text("SET LOCAL app.tenant_id = :tid"), {"tid": tenant_id})
        yield conn
        txn.commit()
    except Exception:
        txn.rollback()
        raise
    finally:
        conn.close()