Agent Zero
126 lines
4.6 KiB
Python
126 lines
4.6 KiB
Python
"""RLS repair + separate DB runtime user.
|
|
|
|
Revision ID: 0044
|
|
Revises: 0043
|
|
Created: 2026-07-26
|
|
|
|
This migration:
|
|
1. Re-discovers ALL tenant-scoped tables and ensures RLS is enabled
|
|
with FORCE + WITH CHECK (covers tables added after migration 0028).
|
|
2. Creates a separate ``crm_runtime`` role with NOSUPERUSER and
|
|
NOBYPASSRLS so the application cannot bypass RLS.
|
|
3. Grants only DML permissions (SELECT/INSERT/UPDATE/DELETE) to
|
|
``crm_runtime`` on all tenant-scoped tables.
|
|
|
|
IMPORTANT: After this migration, the application's DATABASE_URL must
|
|
use ``crm_runtime`` (not the superuser) for API and worker containers.
|
|
Migration/DDL operations continue to use the owner user (crm_user).
|
|
"""
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
revision = "0044"
|
|
down_revision = "0043_backups"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
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]
|
|
|
|
|
|
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 "
|
|
"AND (policyname LIKE '%tenant%' OR policyname LIKE '%isolation%')"
|
|
),
|
|
{"t": table_name},
|
|
)
|
|
return [row[0] for row in result]
|
|
|
|
|
|
def upgrade() -> None:
|
|
conn = op.get_bind()
|
|
|
|
# ── 1. RLS Repair: ensure all tenant tables have RLS + WITH CHECK ──
|
|
tenant_tables = _discover_tenant_tables(conn)
|
|
logger.info("RLS repair: discovered %d tenant-scoped tables: %s", len(tenant_tables), tenant_tables)
|
|
|
|
for table_name in tenant_tables:
|
|
# Enable RLS
|
|
op.execute(f'ALTER TABLE "{table_name}" ENABLE ROW LEVEL SECURITY')
|
|
# Force RLS (applies to table owner too)
|
|
op.execute(f'ALTER TABLE "{table_name}" FORCE ROW LEVEL SECURITY')
|
|
|
|
# Drop existing tenant policies
|
|
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)
|
|
|
|
# Create unified tenant isolation policy with WITH CHECK
|
|
op.execute(
|
|
f'CREATE POLICY tenant_isolation ON "{table_name}" '
|
|
f"USING (tenant_id = current_setting('app.tenant_id', true)::uuid) "
|
|
f"WITH CHECK (tenant_id = current_setting('app.tenant_id', true)::uuid)"
|
|
)
|
|
logger.info("Created/updated tenant_isolation policy on %s (USING + WITH CHECK)", table_name)
|
|
|
|
# ── 2. Create crm_runtime role (NOSUPERUSER, NOBYPASSRLS) ──
|
|
# Use DO block for idempotent creation
|
|
op.execute(
|
|
sa.text(
|
|
"DO $$ "
|
|
"BEGIN "
|
|
" IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'crm_runtime') THEN "
|
|
" CREATE ROLE crm_runtime LOGIN NOSUPERUSER NOCREATEDB NOCREATEROLE "
|
|
" NOREPLICATION NOBYPASSRLS; "
|
|
" END IF; "
|
|
"END $$;"
|
|
)
|
|
)
|
|
logger.info("Ensured crm_runtime role exists (NOSUPERUSER, NOBYPASSRLS)")
|
|
|
|
# ── 3. Grant DML permissions to crm_runtime on all tenant tables ──
|
|
for table_name in tenant_tables:
|
|
op.execute(
|
|
f'GRANT SELECT, INSERT, UPDATE, DELETE ON "{table_name}" TO crm_runtime'
|
|
)
|
|
|
|
# Grant usage on sequences (for SERIAL/IDENTITY columns)
|
|
op.execute("GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO crm_runtime")
|
|
|
|
logger.info("Granted DML permissions to crm_runtime on %d tables", len(tenant_tables))
|
|
|
|
|
|
def downgrade() -> None:
|
|
conn = op.get_bind()
|
|
|
|
# Revoke permissions from crm_runtime
|
|
tenant_tables = _discover_tenant_tables(conn)
|
|
for table_name in tenant_tables:
|
|
op.execute(f'REVOKE SELECT, INSERT, UPDATE, DELETE ON "{table_name}" FROM crm_runtime')
|
|
op.execute("REVOKE USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public FROM crm_runtime")
|
|
|
|
# Drop crm_runtime role
|
|
op.execute("DROP ROLE IF EXISTS crm_runtime")
|
|
logger.info("Dropped crm_runtime role")
|
|
|
|
# Note: RLS policies are NOT reverted here to avoid weakening security.
|
|
# Migration 0028's downgrade handles the original set of tables.
|