Agent Zero
99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
"""Add owner_id column to Phase 2 tables for row-level ownership.
|
|
|
|
Adds nullable owner_id (FK -> users.id, ON DELETE SET NULL) to tables
|
|
that gained OwnedMixin in Phase 2. For tables that already have a
|
|
non-nullable user_id column, owner_id is backfilled from user_id.
|
|
|
|
Plugin tables may not exist yet at migration time (created by plugin
|
|
migrations separately), so we check existence before adding columns.
|
|
|
|
Revision ID: 0102
|
|
"""
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects.postgresql import UUID as PGUUID
|
|
|
|
revision = "0102"
|
|
down_revision = "0101"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
# (table_name, has_user_id_to_backfill)
|
|
TABLES = [
|
|
# Core models (always exist at this migration point)
|
|
("contact_folders", True),
|
|
("user_preferences", True),
|
|
("workspaces", False),
|
|
# Plugin models (may not exist yet — created by plugin migrations)
|
|
("mcp_server_configs", False),
|
|
("automation_agent_definitions", False),
|
|
("automation_definitions", False),
|
|
("report_templates", False),
|
|
("report_instances", False),
|
|
("entity_links", False),
|
|
("comm_conversations", False),
|
|
("ai_proactive_suggestions", True),
|
|
("ai_agents", False),
|
|
("ai_chat_sessions", True),
|
|
("tags", False),
|
|
("share_links", True),
|
|
]
|
|
|
|
|
|
def _table_exists(conn, table_name: str) -> bool:
|
|
result = conn.execute(
|
|
sa.text("SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = :name)"),
|
|
{"name": table_name},
|
|
)
|
|
return result.scalar()
|
|
|
|
|
|
def upgrade() -> None:
|
|
conn = op.get_bind()
|
|
for table_name, has_user_id in TABLES:
|
|
if not _table_exists(conn, table_name):
|
|
print(f"[0102] Skipping {table_name} — table does not exist yet")
|
|
continue
|
|
# Check if owner_id column already exists
|
|
col_exists = conn.execute(
|
|
sa.text(
|
|
"SELECT EXISTS (SELECT 1 FROM information_schema.columns "
|
|
"WHERE table_name = :name AND column_name = 'owner_id')"
|
|
),
|
|
{"name": table_name},
|
|
).scalar()
|
|
if col_exists:
|
|
print(f"[0102] Skipping {table_name} — owner_id already exists")
|
|
continue
|
|
op.add_column(
|
|
table_name,
|
|
sa.Column(
|
|
"owner_id",
|
|
PGUUID(as_uuid=True),
|
|
sa.ForeignKey("users.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
),
|
|
)
|
|
op.create_index(
|
|
f"ix_{table_name}_owner_id",
|
|
table_name,
|
|
["owner_id"],
|
|
)
|
|
if has_user_id:
|
|
op.execute(
|
|
f"UPDATE {table_name} SET owner_id = user_id WHERE owner_id IS NULL;"
|
|
)
|
|
print(f"[0102] Added owner_id to {table_name} (backfilled from user_id)")
|
|
else:
|
|
print(f"[0102] Added owner_id to {table_name}")
|
|
|
|
|
|
def downgrade() -> None:
|
|
conn = op.get_bind()
|
|
for table_name, _ in TABLES:
|
|
if not _table_exists(conn, table_name):
|
|
continue
|
|
op.drop_index(f"ix_{table_name}_owner_id", table_name=table_name)
|
|
op.drop_column(table_name, "owner_id")
|