Source code for codegen_database.ext.rls.alembic

"""Alembic autogenerate integration for RLS policies.

Call :func:`register_rls_alembic_events` in ``env.py`` alongside
:func:`~codegen_database.alembic.register.alembic_hook` to enable
RLS policy diffing in Alembic autogenerate::

    from codegen_database.alembic.register import alembic_hook
    from codegen_database.ext.rls.alembic import register_rls_alembic_events

    alembic_hook()
    register_rls_alembic_events()

The comparator queries ``pg_policies`` and emits ``op.execute(...)``
statements for any policies that need to be created or dropped.
"""

from typing import TYPE_CHECKING

from alembic.autogenerate.compare import comparators
from alembic.autogenerate.render import renderers

from codegen_database.ext.rls.base import RLSPolicies
from codegen_database.ext.rls.compare import (
    CreateRLSPolicyOp,
    DropRLSPolicyOp,
    compare_rls_policies,
)

if TYPE_CHECKING:
    from alembic.autogenerate.api import AutogenContext
    from alembic.operations.ops import UpgradeOps


def _compare_rls(
    autogen_context: AutogenContext,
    upgrade_ops: UpgradeOps,
    _schemas: object,
) -> None:
    """Compare desired RLS policies against the current database state."""
    policies = RLSPolicies.extract(autogen_context.metadata)

    if not policies:
        return

    assert autogen_context.connection  # noqa: S101
    result = compare_rls_policies(autogen_context.connection, policies)
    upgrade_ops.ops.extend(result)


def _render_rls_op(
    _autogen_context: AutogenContext,
    op: CreateRLSPolicyOp | DropRLSPolicyOp,
) -> list[str]:
    """Render an RLS op as ``op.execute(...)`` statements."""
    return [f'op.execute("""{cmd}""")' for cmd in op.to_sql()]




def register_rls_alembic_events() -> None:
    """Register RLS comparator and renderer into Alembic.

    Call once in ``env.py`` before ``context.configure()``.

    """
    comparators.dispatch_for("schema")(_compare_rls)
    renderers.dispatch_for(CreateRLSPolicyOp)(_render_rls_op)
    renderers.dispatch_for(DropRLSPolicyOp)(_render_rls_op)