leocrm/alembic/versions/0111_enable_rls_remaining_tenant_tables.py

"""Enable RLS for all remaining tenant tables that still lack RLS after 0108/0109.

Migration 0108 dynamically discovered tables with tenant_id and enabled RLS.
However, new tables may have been added since, or some were missed.

This migration re-runs the same dynamic discovery to catch any stragglers.

Login tables (users, user_tenants, tenants, sessions, password_reset_tokens,
roles, permissions) are explicitly excluded — they must NOT have RLS.

Global tables (alembic_version, plugin_migrations, marketplace_listings,
sequences, notification_types) are also excluded.

Revision ID: 0111
Revises: 0110
"""

from __future__ import annotations

from alembic import op

revision = "0111"
down_revision = "0110"
branch_labels = None
depends_on = None


# Tables that must never get RLS (global / cross-tenant infrastructure)
GLOBAL_TABLES = [
    "alembic_version",
    "plugin_migrations",
    "marketplace_listings",
    "sequences",
    "notification_types",
]

# Login-related tables — RLS blocks crm_auth during login flow
# See 0108 and 0109 for detailed explanation
LOGIN_TABLES = [
    "users",
    "user_tenants",
    "tenants",
    "sessions",
    "password_reset_tokens",
    "roles",
    "permissions",
]


def _exec(sql: str) -> None:
    op.execute(sql)


def upgrade() -> None:
    # Dynamic discovery + RLS activation for any tenant table still missing RLS
    _exec("""
DO $$
DECLARE
    r RECORD;
    policy_sql TEXT;
BEGIN
    FOR r IN
        SELECT t.table_name
        FROM information_schema.tables t
        JOIN information_schema.columns c
          ON c.table_schema = t.table_schema
         AND c.table_name = t.table_name
         AND c.column_name = 'tenant_id'
        WHERE t.table_schema = 'public'
          AND t.table_type = 'BASE TABLE'
          AND t.table_name NOT IN (
              'alembic_version',
              'plugin_migrations',
              'marketplace_listings',
              'sequences',
              'notification_types',
              -- Login tables must NOT have RLS
              'users',
              'user_tenants',
              'tenants',
              'sessions',
              'password_reset_tokens',
              'roles',
              'permissions'
          )
          AND NOT EXISTS (
              SELECT 1
              FROM pg_class pc
              JOIN pg_namespace pn ON pn.oid = pc.relnamespace
              WHERE pn.nspname = 'public'
                AND pc.relname = t.table_name
                AND pc.relrowsecurity = true
          )
    LOOP
        -- Enable + force RLS
        EXECUTE format('ALTER TABLE public.%I ENABLE ROW LEVEL SECURITY', r.table_name);
        EXECUTE format('ALTER TABLE public.%I FORCE ROW LEVEL SECURITY', r.table_name);

        -- Drop stale policies (idempotent)
        EXECUTE format('DROP POLICY IF EXISTS tenant_isolation ON public.%I', r.table_name);
        EXECUTE format('DROP POLICY IF EXISTS %s_tenant_isolation ON public.%I', r.table_name, r.table_name);

        -- Create fail-closed policy
        policy_sql := format(
            'CREATE POLICY %s_tenant_isolation '
            'ON public.%I '
            'FOR ALL '
            'TO crm_api, crm_worker '
            'USING (tenant_id = NULLIF(current_setting(''app.current_tenant_id'', true), '''')::uuid) '
            'WITH CHECK (tenant_id = NULLIF(current_setting(''app.current_tenant_id'', true), '''')::uuid)',
            r.table_name, r.table_name
        );
        EXECUTE policy_sql;

        -- Grant CRUD to crm_api and crm_worker
        EXECUTE format('GRANT SELECT, INSERT, UPDATE, DELETE ON public.%I TO crm_api', r.table_name);
        EXECUTE format('GRANT SELECT, INSERT, UPDATE, DELETE ON public.%I TO crm_worker', r.table_name);
    END LOOP;
END $$;
    """)


def downgrade() -> None:
    pass