Source code for codegen_database.ext.rls.base

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from sqlalchemy import MetaData




@dataclass(frozen=True, eq=True)
class RLSPolicy:
    table: str
    schema: str
    name: str
    using: str
    with_check: str | None = None
    roles: list[str] = field(default_factory=list, hash=False)
    permissive: bool = True
    cmd: str = "ALL"

    @property
    def table_fullname(self) -> str:
        return f"{self.schema}.{self.table}"

    def to_sql_create(self) -> str:
        parts = [
            f"CREATE POLICY {self.name}",
            f"ON {self.table_fullname}",
            "AS PERMISSIVE" if self.permissive else "AS RESTRICTIVE",
            f"FOR {self.cmd}",
        ]

        if self.roles:
            roles_str = ", ".join(self.roles)
            parts.append(f"TO {roles_str}")

        parts.append(f"USING ({self.using})")

        if self.with_check is not None:
            parts.append(f"WITH CHECK ({self.with_check})")

        return "\n".join(parts)

    def to_sql_drop(self) -> str:
        return f"DROP POLICY IF EXISTS {self.name} ON {self.table_fullname}"

    def to_sql_enable_rls(self) -> str:
        return f"ALTER TABLE {self.table_fullname} ENABLE ROW LEVEL SECURITY"





@dataclass
class RLSPolicies:
    policies: set[RLSPolicy] = field(default_factory=set)

    @classmethod
    def extract(
        cls,
        metadata: MetaData,
    ) -> RLSPolicies:
        rls_policies = metadata.info.get("rls_policies", RLSPolicies())
        metadata.info["rls_policies"] = rls_policies
        return rls_policies



def register_rls_policy(
    metadata: MetaData,
    policy: RLSPolicy,
) -> None:
    RLSPolicies.extract(metadata).policies.add(policy)