leocrm/alembic/versions/0028_rls_force.py

"""FORCE Row Level Security + WITH CHECK on all tenant-scoped tables.

Revision ID: 0028_rls_force
Revises: 0027_unify_company_to_contact
Create Date: 2026-07-25

This migration:
1. Discovers all tables in the public schema that have a tenant_id column.
2. ALTER TABLE ... FORCE ROW LEVEL SECURITY on each (ensures RLS applies to table owners too).
3. Drops existing tenant_isolation policies and recreates them with both
   USING and WITH CHECK clauses so writes are also filtered by tenant.
4. Covers core tables AND plugin tables (anything with tenant_id).

Idempotent: safe to run multiple times.
"""

from __future__ import annotations

import logging
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa

revision: str = "0028_rls_force"
down_revision: Union[str, None] = "0027_unify_company_to_contact"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

logger = logging.getLogger("alembic.migration.0028_rls_force")


def _discover_tenant_tables(conn) -> list[str]:
    """Return all table names in the public schema that have a tenant_id column."""
    result = conn.execute(
        sa.text(
            "SELECT table_name FROM information_schema.columns "
            "WHERE table_schema = 'public' AND column_name = 'tenant_id' "
            "ORDER BY table_name"
        )
    )
    return [row[0] for row in result.fetchall()]


def _discover_existing_policies(conn, table_name: str) -> list[str]:
    """Return all policy names on *table_name* that contain 'tenant' or 'isolation'."""
    result = conn.execute(
        sa.text(
            "SELECT policyname FROM pg_policies "
            "WHERE schemaname = 'public' AND tablename = :t"
        ),
        {"t": table_name},
    )
    return [row[0] for row in result.fetchall()]


def upgrade() -> None:
    conn = op.get_bind()
    tenant_tables = _discover_tenant_tables(conn)
    logger.info("Discovered %d tenant-scoped tables: %s", len(tenant_tables), tenant_tables)

    for table_name in tenant_tables:
        # 1. Enable RLS (idempotent — ENABLE is safe to repeat)
        op.execute(f'ALTER TABLE "{table_name}" ENABLE ROW LEVEL SECURITY')

        # 2. FORCE RLS — ensures policies apply even to table owners/superusers
        #    who would otherwise bypass RLS
        op.execute(f'ALTER TABLE "{table_name}" FORCE ROW LEVEL SECURITY')

        # 3. Drop ALL existing policies on this table that relate to tenant isolation
        existing_policies = _discover_existing_policies(conn, table_name)
        for policy_name in existing_policies:
            op.execute(f'DROP POLICY IF EXISTS "{policy_name}" ON "{table_name}"')
            logger.info("Dropped policy %s on %s", policy_name, table_name)

        # 4. Create new policy with both USING and WITH CHECK
        #    USING:   filters rows visible in SELECT/UPDATE/DELETE
        #    WITH CHECK: enforces tenant_id on INSERT/UPDATE
        op.execute(
            f'CREATE POLICY tenant_isolation ON "{table_name}" '
            f"USING (tenant_id = current_setting('app.current_tenant_id')::uuid) "
            f"WITH CHECK (tenant_id = current_setting('app.current_tenant_id')::uuid)"
        )
        logger.info("Created policy tenant_isolation on %s (USING + WITH CHECK)", table_name)


def downgrade() -> None:
    """Revert FORCE RLS and restore USING-only policies (matching 0015 behavior)."""
    conn = op.get_bind()
    tenant_tables = _discover_tenant_tables(conn)

    for table_name in tenant_tables:
        # Drop the USING+WITH CHECK policy
        op.execute(f'DROP POLICY IF EXISTS tenant_isolation ON "{table_name}"')

        # Remove FORCE but keep ENABLE (matching pre-0028 state)
        op.execute(f'ALTER TABLE "{table_name}" NO FORCE ROW LEVEL SECURITY')

        # Recreate USING-only policy (matching original 0015 behavior)
        op.execute(
            f'CREATE POLICY tenant_isolation ON "{table_name}" '
            f"USING (tenant_id = current_setting('app.current_tenant_id')::uuid)"
        )
        logger.info("Reverted %s to USING-only policy (removed FORCE, removed WITH CHECK)", table_name)