Agent Zero
64 lines
1.7 KiB
Python
64 lines
1.7 KiB
Python
"""Add owner_id to plugin entity tables for row-level ownership.
|
|
|
|
Revision ID: 0054
|
|
Revises: 0053
|
|
Create Date: 2026-07-29
|
|
"""
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects.postgresql import UUID as PGUUID
|
|
|
|
revision = "0054"
|
|
down_revision = "0053"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
# Tables that need owner_id
|
|
TABLES = [
|
|
"files",
|
|
"folders",
|
|
"calendar_entries",
|
|
"calendars",
|
|
"tasks",
|
|
"subtasks",
|
|
]
|
|
|
|
|
|
def upgrade() -> None:
|
|
# Check which columns already exist before adding
|
|
conn = op.get_bind()
|
|
for table in TABLES:
|
|
# Check if table exists
|
|
table_exists = conn.execute(
|
|
sa.text(
|
|
"SELECT 1 FROM information_schema.tables WHERE table_name = :table"
|
|
),
|
|
{"table": table},
|
|
).fetchone()
|
|
if table_exists is None:
|
|
continue
|
|
# Check if column already exists
|
|
result = conn.execute(
|
|
sa.text(
|
|
"SELECT column_name FROM information_schema.columns "
|
|
"WHERE table_name = :table AND column_name = 'owner_id'"
|
|
),
|
|
{"table": table},
|
|
)
|
|
if result.fetchone() is None:
|
|
op.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS owner_id UUID REFERENCES users(id) ON DELETE SET NULL")
|
|
op.execute(f"CREATE INDEX IF NOT EXISTS ix_{table}_owner ON {table} (owner_id)")
|
|
|
|
|
|
def downgrade() -> None:
|
|
for table in TABLES:
|
|
try:
|
|
op.drop_index(f"ix_{table}_owner", table_name=table)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
op.drop_column(table, "owner_id")
|
|
except Exception:
|
|
pass
|