Agent Zero
63 lines
1.5 KiB
Python
63 lines
1.5 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 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.add_column(
|
|
table,
|
|
sa.Column(
|
|
"owner_id",
|
|
PGUUID(as_uuid=True),
|
|
sa.ForeignKey("users.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
),
|
|
)
|
|
op.create_index(f"ix_{table}_owner", 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
|