Agent Zero
69 lines
2.8 KiB
Python
69 lines
2.8 KiB
Python
"""Enable RLS on critical tables missing it.
|
|
|
|
Tables: api_tokens, sequences, tenant_plugin_activation, user_tenants.
|
|
Also adds tenant_id to guest_invitations and enables RLS.
|
|
|
|
NOTE: sessions and password_reset_tokens are EXCLUDED from RLS because
|
|
crm_auth accesses them without a tenant context (during login/reset).
|
|
Instead, their security is enforced via GRANT restrictions in migration 0100.
|
|
|
|
Revision ID: 0101
|
|
"""
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
revision = "0101"
|
|
down_revision = "0100"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
# Tables that have tenant_id and can safely get RLS
|
|
# (accessed only by crm_api/crm_worker which always set tenant context)
|
|
RLS_TABLES = [
|
|
"api_tokens",
|
|
"sequences",
|
|
"tenant_plugin_activation",
|
|
"user_tenants",
|
|
]
|
|
|
|
|
|
def upgrade() -> None:
|
|
# Undo any manual RLS changes on auth tables (safety measure)
|
|
op.execute("ALTER TABLE password_reset_tokens DISABLE ROW LEVEL SECURITY;")
|
|
op.execute("ALTER TABLE sessions DISABLE ROW LEVEL SECURITY;")
|
|
op.execute("DROP POLICY IF EXISTS password_reset_tokens_tenant_isolation ON password_reset_tokens;")
|
|
op.execute("DROP POLICY IF EXISTS sessions_tenant_isolation ON sessions;")
|
|
|
|
# Enable RLS + create tenant isolation policy for each table
|
|
for table in RLS_TABLES:
|
|
op.execute(f"ALTER TABLE {table} ENABLE ROW LEVEL SECURITY;")
|
|
op.execute(
|
|
f"CREATE POLICY {table}_tenant_isolation ON {table} "
|
|
f"FOR ALL USING (tenant_id = NULLIF(current_setting('app.current_tenant_id', true), '')::uuid) "
|
|
f"WITH CHECK (tenant_id = NULLIF(current_setting('app.current_tenant_id', true), '')::uuid);"
|
|
)
|
|
|
|
# guest_invitations: add tenant_id + enable RLS
|
|
op.add_column("guest_invitations", sa.Column("tenant_id", sa.UUID(), nullable=True))
|
|
op.execute("CREATE INDEX ix_guest_invitations_tenant_id ON guest_invitations (tenant_id);")
|
|
op.execute("ALTER TABLE guest_invitations ENABLE ROW LEVEL SECURITY;")
|
|
op.execute(
|
|
"CREATE POLICY guest_invitations_tenant_isolation ON guest_invitations "
|
|
"FOR ALL 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);"
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
# Drop guest_invitations RLS + tenant_id
|
|
op.execute("DROP POLICY IF EXISTS guest_invitations_tenant_isolation ON guest_invitations;")
|
|
op.execute("ALTER TABLE guest_invitations DISABLE ROW LEVEL SECURITY;")
|
|
op.execute("DROP INDEX IF EXISTS ix_guest_invitations_tenant_id;")
|
|
op.drop_column("guest_invitations", "tenant_id")
|
|
|
|
# Drop RLS on other tables
|
|
for table in RLS_TABLES:
|
|
op.execute(f"DROP POLICY IF EXISTS {table}_tenant_isolation ON {table};")
|
|
op.execute(f"ALTER TABLE {table} DISABLE ROW LEVEL SECURITY;")
|