leocrm/alembic/versions/0091_add_tenant_fk_constraints.py

"""Add tenant_id FK constraints to all tenant-scoped tables.

Phase 2 Data Integrity: Adds FOREIGN KEY (tenant_id) REFERENCES tenants(id)
ON DELETE CASCADE to all tenant-scoped tables that have a tenant_id column
but no FK constraint yet.

Global tables (sequences, system_settings, currencies, tax_rates, permissions,
permission_templates, unified_search_providers, unified_search_index_log,
mcp_server_configs, plugin_test_data) are excluded because they use tenant_id
for filtering but are not owned by a single tenant.

Revision ID: 0091
Revises: 0090
"""

from __future__ import annotations

from alembic import op
import sqlalchemy as sa

revision = "0091"
down_revision = "0090"
branch_labels = None
depends_on = None

# All 74 tenant-scoped tables that need FK constraints.
# Excludes 10 global tables that use tenant_id but are not tenant-owned.
TENANT_TABLES = [
    "addresses",
    "ai_agents",
    "ai_chat_attachments",
    "ai_chat_folders",
    "ai_chat_messages",
    "ai_chat_sessions",
    "ai_models",
    "ai_presets",
    "ai_proactive_context_log",
    "ai_proactive_settings",
    "ai_proactive_suggestions",
    "ai_providers",
    "attachments",
    "automation_agent_definitions",
    "automation_agent_runs",
    "automation_agent_versions",
    "automation_cron_jobs",
    "automation_definitions",
    "automation_runs",
    "automation_versions",
    "backups",
    "bank_accounts",
    "calendar_entries",
    "calendar_entry_links",
    "calendar_shares",
    "calendars",
    "comm_conversation_mutes",
    "comm_conversation_pins",
    "comm_conversations",
    "comm_message_attachments",
    "comm_message_blocks",
    "comm_message_edits",
    "comm_message_reactions",
    "comm_message_reads",
    "comm_messages",
    "comm_participants",
    "contact_folders",
    "contact_pgp_keys",
    "contacts",
    "custom_field_definitions",
    "entity_links",
    "entity_policies",
    "event_outbox",
    "files",
    "folders",
    "mail_account_delegates",
    "mail_account_send_permissions",
    "mail_accounts",
    "mail_attachments",
    "mail_folders",
    "mail_label_assignments",
    "mail_labels",
    "mail_rules",
    "mail_seen_by",
    "mail_signatures",
    "mail_sync_queue",
    "mail_templates",
    "mails",
    "permission_delegations",
    "pgp_keys",
    "report_instances",
    "report_templates",
    "resource_bookings",
    "resources",
    "saved_filters",
    "saved_views",
    "share_links",
    "subtasks",
    "tag_assignments",
    "tags",
    "tasks",
    "user_calendar_visibility",
    "vacation_sent_log",
    "webhooks",
]


def upgrade() -> None:
    # Step 1: Clean orphaned tenant_id references before adding FK constraints.
    # Set tenant_id = NULL where the referenced tenant does not exist.
    for table in TENANT_TABLES:
        op.execute(f"""
            DO $$ BEGIN
                IF EXISTS (
                    SELECT 1 FROM information_schema.tables
                    WHERE table_schema = 'public' AND table_name = '{table}'
                ) AND EXISTS (
                    SELECT 1 FROM information_schema.columns
                    WHERE table_schema = 'public' AND table_name = '{table}'
                      AND column_name = 'tenant_id'
                ) THEN
                    UPDATE public.{table}
                    SET tenant_id = NULL
                    WHERE tenant_id IS NOT NULL
                      AND tenant_id NOT IN (SELECT id FROM public.tenants);
                END IF;
            END $$;
        """)

    # Step 2: Add FK constraints idempotently.
    for table in TENANT_TABLES:
        constraint_name = f"fk_{table}_tenant_id"
        op.execute(f"""
            DO $$ BEGIN
                IF EXISTS (
                    SELECT 1 FROM information_schema.tables
                    WHERE table_schema = 'public' AND table_name = '{table}'
                ) AND EXISTS (
                    SELECT 1 FROM information_schema.columns
                    WHERE table_schema = 'public' AND table_name = '{table}'
                      AND column_name = 'tenant_id'
                ) AND NOT EXISTS (
                    SELECT 1 FROM information_schema.table_constraints
                    WHERE constraint_schema = 'public'
                      AND constraint_name = '{constraint_name}'
                      AND constraint_type = 'FOREIGN KEY'
                ) THEN
                    ALTER TABLE public.{table}
                    ADD CONSTRAINT {constraint_name}
                    FOREIGN KEY (tenant_id)
                    REFERENCES public.tenants(id)
                    ON DELETE CASCADE;
                END IF;
            END $$;
        """)


def downgrade() -> None:
    for table in TENANT_TABLES:
        constraint_name = f"fk_{table}_tenant_id"
        op.execute(f"""
            DO $$ BEGIN
                IF EXISTS (
                    SELECT 1 FROM information_schema.table_constraints
                    WHERE constraint_schema = 'public'
                      AND constraint_name = '{constraint_name}'
                      AND constraint_type = 'FOREIGN KEY'
                ) THEN
                    ALTER TABLE public.{table}
                    DROP CONSTRAINT {constraint_name};
                END IF;
            END $$;
        """)