leocrm/alembic/versions/0102_add_owner_id_to_phase2_tables.py

"""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")