Agent Zero
61 lines
2.1 KiB
Python
61 lines
2.1 KiB
Python
"""Enable RLS on critical tables missing it.
|
|
|
|
Tables: api_tokens, sequences, sessions, tenant_plugin_activation,
|
|
user_tenants, password_reset_tokens.
|
|
Also adds tenant_id to guest_invitations and enables RLS.
|
|
|
|
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 but no RLS
|
|
RLS_TABLES = [
|
|
"api_tokens",
|
|
"sequences",
|
|
"sessions",
|
|
"tenant_plugin_activation",
|
|
"user_tenants",
|
|
"password_reset_tokens",
|
|
]
|
|
|
|
|
|
def upgrade() -> None:
|
|
# 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 = current_setting('app.current_tenant_id')::uuid) "
|
|
f"WITH CHECK (tenant_id = current_setting('app.current_tenant_id')::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 = current_setting('app.current_tenant_id')::uuid) "
|
|
"WITH CHECK (tenant_id = current_setting('app.current_tenant_id')::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;")
|