"""Row-level security plugins for codegen_database dimensions.
Each plugin registers one or more PostgreSQL RLS policies on the raw
backing table so that queries are automatically filtered by the
current session's ``app.*`` settings.
Four pre-built policies cover the most common patterns:
- :class:`OwnerRLSPlugin` -- generic: any column matched against any
``app.*`` setting.
- :class:`UserRLSPlugin` -- per-user isolation (``app.user_id``).
- :class:`TenantRLSPlugin` -- per-tenant isolation
(``app.tenant_id``).
- :class:`TenantUserRLSPlugin` -- tenant isolation AND per-user
visibility within the tenant.
All plugins accept an optional *bypass_roles* list. Roles in that
list receive an unconditional ``USING (true)`` policy so they can
read and write every row -- useful for service accounts and trusted
roles.
Prerequisites
-------------
1. RLS policy diffing in Alembic requires calling
:func:`~codegen_database.ext.rls.alembic.register_rls_alembic_events` once
in ``env.py``.
2. The application must call
``SET LOCAL app.user_id = '...'`` (or equivalent) for each
request so the session setting is available.
Usage::
class Users(Base):
__tablename__ = "users"
__table_args__ = {"schema": "public"}
__plugins__ = [UserRLSPlugin(bypass_roles=["authenticator"])]
user_id = Column(String, nullable=False)
name = Column(String, nullable=False)
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from codegen_database.ext.rls.base import RLSPolicy, register_rls_policy
from codegen_database.plugin import Dynamic, Plugin, requires
if TYPE_CHECKING:
from codegen_database.factory.context import FactoryContext
def _bypass_policy(table: str, schema: str, role: str) -> RLSPolicy:
"""Return an unconditional bypass policy for *role*.
Args:
table: Unqualified table name.
schema: Schema name.
role: PostgreSQL role name.
Returns:
An :class:`~codegen_database.ext.rls.base.RLSPolicy` that allows *role*
to read and write every row unconditionally.
"""
safe_role = role.replace("-", "_")
return RLSPolicy(
table=table,
schema=schema,
name=f"bypass_{safe_role}",
using="true",
with_check="true",
roles=[role],
)
[docs]
@requires(Dynamic("table_key"))
class OwnerRLSPlugin(Plugin):
"""Generic owner-column RLS: filter rows by a session setting.
Enables row-level security on the target table and creates a
``PERMISSIVE`` policy that restricts access to rows where
``current_setting(setting, true)::text = column::text``.
Args:
column: Column that identifies the row's owner.
setting: PostgreSQL session setting to compare against,
e.g. ``"app.user_id"``.
policy_name: Name for the created policy
(default ``"owner_policy"``).
bypass_roles: Roles that receive an unconditional
``USING (true)`` policy, bypassing the owner filter.
table_key: Key in ``ctx`` for the table to protect
(default ``"raw_table"``).
"""
def __init__(
self,
column: str,
setting: str,
*,
policy_name: str = "owner_policy",
bypass_roles: list[str] | None = None,
table_key: str = "raw_table",
) -> None:
"""Store configuration."""
self.column = column
self.setting = setting
self.policy_name = policy_name
self.bypass_roles = bypass_roles or []
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Register the owner RLS policy on the backing table."""
table = ctx[self.table_key]
tname = table.name
schema = ctx.schemaname
using = (
f"current_setting({self.setting!r}, true)::text"
f" = {self.column}::text"
)
register_rls_policy(
ctx.metadata,
RLSPolicy(
table=tname,
schema=schema,
name=self.policy_name,
using=using,
with_check=using,
),
)
for role in self.bypass_roles:
register_rls_policy(
ctx.metadata,
_bypass_policy(tname, schema, role),
)
[docs]
@requires(Dynamic("table_key"))
class UserRLSPlugin(Plugin):
"""Per-user row isolation via ``app.user_id``.
Filters rows to those where the nominated *user_column* matches
``current_setting('app.user_id', true)``.
Args:
user_column: Column holding the owning user ID
(default ``"user_id"``).
bypass_roles: Roles that bypass the filter.
table_key: Key in ``ctx`` for the backing table
(default ``"raw_table"``).
"""
def __init__(
self,
user_column: str = "user_id",
*,
bypass_roles: list[str] | None = None,
table_key: str = "raw_table",
) -> None:
"""Store configuration."""
self.user_column = user_column
self.bypass_roles = bypass_roles or []
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Register user-isolation RLS policy."""
table = ctx[self.table_key]
tname = table.name
schema = ctx.schemaname
using = (
"current_setting('app.user_id', true)::text"
f" = {self.user_column}::text"
)
register_rls_policy(
ctx.metadata,
RLSPolicy(
table=tname,
schema=schema,
name="user_policy",
using=using,
with_check=using,
),
)
for role in self.bypass_roles:
register_rls_policy(
ctx.metadata,
_bypass_policy(tname, schema, role),
)
[docs]
@requires(Dynamic("table_key"))
class TenantRLSPlugin(Plugin):
"""Per-tenant row isolation via ``app.tenant_id``.
Filters rows to those where the nominated *tenant_column* matches
``current_setting('app.tenant_id', true)``.
Args:
tenant_column: Column holding the tenant ID
(default ``"tenant_id"``).
bypass_roles: Roles that bypass the filter.
table_key: Key in ``ctx`` for the backing table
(default ``"raw_table"``).
"""
def __init__(
self,
tenant_column: str = "tenant_id",
*,
bypass_roles: list[str] | None = None,
table_key: str = "raw_table",
) -> None:
"""Store configuration."""
self.tenant_column = tenant_column
self.bypass_roles = bypass_roles or []
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Register tenant-isolation RLS policy."""
table = ctx[self.table_key]
tname = table.name
schema = ctx.schemaname
using = (
"current_setting('app.tenant_id', true)::text"
f" = {self.tenant_column}::text"
)
register_rls_policy(
ctx.metadata,
RLSPolicy(
table=tname,
schema=schema,
name="tenant_policy",
using=using,
with_check=using,
),
)
for role in self.bypass_roles:
register_rls_policy(
ctx.metadata,
_bypass_policy(tname, schema, role),
)
[docs]
@requires(Dynamic("table_key"))
class TenantUserRLSPlugin(Plugin):
"""Tenant isolation combined with per-user visibility.
Rows are only visible when BOTH conditions hold:
- ``current_setting('app.tenant_id', true)`` matches
*tenant_column*, AND
- ``current_setting('app.user_id', true)`` matches *user_column*.
Args:
tenant_column: Column holding the tenant ID
(default ``"tenant_id"``).
user_column: Column holding the user ID
(default ``"user_id"``).
bypass_roles: Roles that bypass the filter.
table_key: Key in ``ctx`` for the backing table
(default ``"raw_table"``).
"""
def __init__(
self,
tenant_column: str = "tenant_id",
user_column: str = "user_id",
*,
bypass_roles: list[str] | None = None,
table_key: str = "raw_table",
) -> None:
"""Store configuration."""
self.tenant_column = tenant_column
self.user_column = user_column
self.bypass_roles = bypass_roles or []
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Register combined tenant+user RLS policy."""
table = ctx[self.table_key]
tname = table.name
schema = ctx.schemaname
tenant_using = (
"current_setting('app.tenant_id', true)::text"
f" = {self.tenant_column}::text"
)
user_using = (
"current_setting('app.user_id', true)::text"
f" = {self.user_column}::text"
)
combined = f"({tenant_using}) AND ({user_using})"
register_rls_policy(
ctx.metadata,
RLSPolicy(
table=tname,
schema=schema,
name="tenant_user_policy",
using=combined,
with_check=combined,
),
)
for role in self.bypass_roles:
register_rls_policy(
ctx.metadata,
_bypass_policy(tname, schema, role),
)