Agent Zero
54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
"""Disable RLS on automation tables (written at startup without tenant context).
|
|
|
|
The automation plugin registers cron jobs and definitions during plugin
|
|
activation, which happens at startup before a tenant context is set.
|
|
RLS blocks these INSERTs because app.current_tenant_id is a dummy default.
|
|
|
|
Revision ID: 0078
|
|
Revises: 0077
|
|
"""
|
|
from alembic import op
|
|
from sqlalchemy import text
|
|
|
|
revision = "0078"
|
|
down_revision = "0077"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
TABLES = [
|
|
"automation_agent_definitions",
|
|
"automation_agent_runs",
|
|
"automation_agent_versions",
|
|
"automation_cron_jobs",
|
|
"automation_definitions",
|
|
"automation_runs",
|
|
"automation_versions",
|
|
]
|
|
|
|
def upgrade() -> None:
|
|
conn = op.get_bind()
|
|
for table in TABLES:
|
|
exists = conn.execute(
|
|
text(f"SELECT 1 FROM information_schema.tables WHERE table_name = '{table}'")
|
|
).fetchone() is not None
|
|
if not exists:
|
|
continue
|
|
conn.execute(text(f"DROP POLICY IF EXISTS {table}_tenant_isolation ON {table}"))
|
|
conn.execute(text(f"ALTER TABLE {table} DISABLE ROW LEVEL SECURITY"))
|
|
|
|
def downgrade() -> None:
|
|
conn = op.get_bind()
|
|
for table in TABLES:
|
|
exists = conn.execute(
|
|
text(f"SELECT 1 FROM information_schema.tables WHERE table_name = '{table}'")
|
|
).fetchone() is not None
|
|
if not exists:
|
|
continue
|
|
conn.execute(text(f"ALTER TABLE {table} ENABLE ROW LEVEL SECURITY"))
|
|
conn.execute(text(
|
|
f"CREATE POLICY {table}_tenant_isolation ON {table} "
|
|
"FOR ALL "
|
|
"USING (tenant_id = current_setting('app.current_tenant_id', true)::uuid) "
|
|
"WITH CHECK (tenant_id = current_setting('app.current_tenant_id', true)::uuid)"
|
|
))
|