leocrm/alembic/versions/0027_unify_company_to_contact.py

"""Unify entity_type 'company' to 'contact' across all plugins.

Revision ID: 0027
Revises: 0026_mail_salt_security
Create Date: 2026-07-23

SAFE MIGRATION: When both company_id and contact_id columns exist in mails,
company_id values are copied to contact_id (where contact_id IS NULL) before
the column is dropped. A backup column is created to track which rows were
originally linked to companies for safe downgrade.

Changes:
- UPDATE entity_links SET entity_type='contact' WHERE entity_type='company'
- UPDATE tag_assignments SET entity_type='contact' WHERE entity_type='company'
- UPDATE calendar_entry_links SET entity_type='contact' WHERE entity_type='company'
- UPDATE addresses SET entity_type='contact' WHERE entity_type='company'
- mails: copy company_id → contact_id WHERE contact_id IS NULL, then drop company_id
"""

from __future__ import annotations

import logging
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa

revision: str = "0027_unify_company_to_contact"
down_revision: Union[str, None] = "0026_mail_salt_security"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

logger = logging.getLogger("alembic.migration.0027")


def _column_exists(conn, table_name: str, column_name: str) -> bool:
    """Check whether *column_name* exists on *table_name* in public schema."""
    result = conn.execute(
        sa.text(
            "SELECT 1 FROM information_schema.columns "
            "WHERE table_schema = 'public' "
            "AND table_name = :t AND column_name = :c"
        ),
        {"t": table_name, "c": column_name},
    ).fetchone()
    return result is not None


def _table_exists(conn, table_name: str) -> bool:
    """Check whether *table_name* exists in public schema."""
    result = conn.execute(
        sa.text(
            "SELECT 1 FROM information_schema.tables "
            "WHERE table_schema = 'public' AND table_name = :t"
        ),
        {"t": table_name},
    ).fetchone()
    return result is not None


def upgrade() -> None:
    conn = op.get_bind()

    # ── 1. Update entity_type: 'company' → 'contact' across link tables ──

    if _table_exists(conn, "entity_links"):
        result = conn.execute(
            sa.text("UPDATE entity_links SET entity_type = 'contact' WHERE entity_type = 'company'")
        )
        logger.info("Updated %d rows in entity_links (company → contact)", result.rowcount)

    if _table_exists(conn, "tag_assignments"):
        result = conn.execute(
            sa.text("UPDATE tag_assignments SET entity_type = 'contact' WHERE entity_type = 'company'")
        )
        logger.info("Updated %d rows in tag_assignments (company → contact)", result.rowcount)

    if _table_exists(conn, "calendar_entry_links"):
        result = conn.execute(
            sa.text("UPDATE calendar_entry_links SET entity_type = 'contact' WHERE entity_type = 'company'")
        )
        logger.info("Updated %d rows in calendar_entry_links (company → contact)", result.rowcount)

    if _table_exists(conn, "addresses"):
        result = conn.execute(
            sa.text("UPDATE addresses SET entity_type = 'contact' WHERE entity_type = 'company'")
        )
        logger.info("Updated %d rows in addresses (company → contact)", result.rowcount)

    # ── 2. Mails: unify company_id into contact_id ──────────────────────
    if not _table_exists(conn, "mails"):
        logger.info("Table 'mails' does not exist — skipping column migration")
        return

    has_company_id = _column_exists(conn, "mails", "company_id")
    has_contact_id = _column_exists(conn, "mails", "contact_id")

    if has_company_id and has_contact_id:
        # Both columns exist: copy company_id → contact_id WHERE contact_id IS NULL
        result = conn.execute(
            sa.text(
                "UPDATE mails SET contact_id = company_id "
                "WHERE contact_id IS NULL AND company_id IS NOT NULL"
            )
        )
        logger.info("Copied %d rows from company_id → contact_id in mails", result.rowcount)

        # Create a backup marker column to track rows originally linked via company_id
        # This enables a targeted downgrade (only revert these rows, not all contact rows)
        if not _column_exists(conn, "mails", "_orig_company_id"):
            op.add_column("mails", sa.Column("_orig_company_id", sa.dialects.postgresql.UUID(as_uuid=True), nullable=True))
            # Record which rows had company_id set (these came from companies)
            op.execute(
                "UPDATE mails SET _orig_company_id = company_id WHERE company_id IS NOT NULL"
            )
            logger.info("Created _orig_company_id backup column for downgrade tracking")

        # Now safe to drop company_id
        op.drop_column("mails", "company_id")
        logger.info("Dropped column company_id from mails")

    elif has_company_id and not has_contact_id:
        # Only company_id exists: simple rename
        op.alter_column("mails", "company_id", new_column_name="contact_id")
        logger.info("Renamed company_id → contact_id in mails")

    else:
        logger.info("No company_id column in mails — nothing to do")


def downgrade() -> None:
    conn = op.get_bind()

    # ── 1. Revert mails: contact_id → company_id ────────────────────────
    if not _table_exists(conn, "mails"):
        return

    has_contact_id = _column_exists(conn, "mails", "contact_id")
    has_company_id = _column_exists(conn, "mails", "company_id")
    has_orig = _column_exists(conn, "mails", "_orig_company_id")

    if has_contact_id and not has_company_id:
        if has_orig:
            # Targeted revert: only restore rows that originally came from company_id
            # Re-add company_id column
            op.add_column("mails", sa.Column("company_id", sa.dialects.postgresql.UUID(as_uuid=True), nullable=True))
            # Restore company_id from the backup marker where it was originally set
            op.execute(
                "UPDATE mails SET company_id = _orig_company_id WHERE _orig_company_id IS NOT NULL"
            )
            # Clear contact_id for rows that were originally company links
            # (only where contact_id matches the original company_id, i.e. it was copied)
            op.execute(
                "UPDATE mails SET contact_id = NULL "
                "WHERE _orig_company_id IS NOT NULL AND contact_id = _orig_company_id"
            )
            # Drop the backup marker
            op.drop_column("mails", "_orig_company_id")
            logger.info("Restored company_id from _orig_company_id backup (targeted revert)")
        else:
            # No backup column — simple rename (fallback for clean installs)
            op.alter_column("mails", "contact_id", new_column_name="company_id")
            logger.info("Renamed contact_id → company_id in mails (no backup marker)")

    # ── 2. Revert entity_type: 'contact' → 'company' ────────────────────
    # NOTE: This is a lossy revert — we cannot distinguish rows that were
    # originally 'company' from rows that were always 'contact'. This only
    # reverts rows that are currently 'contact' back to 'company'.
    # A proper revert requires application-level audit logs.

    if _table_exists(conn, "entity_links"):
        conn.execute(
            sa.text("UPDATE entity_links SET entity_type = 'company' WHERE entity_type = 'contact'")
        )
    if _table_exists(conn, "tag_assignments"):
        conn.execute(
            sa.text("UPDATE tag_assignments SET entity_type = 'company' WHERE entity_type = 'contact'")
        )
    if _table_exists(conn, "calendar_entry_links"):
        conn.execute(
            sa.text("UPDATE calendar_entry_links SET entity_type = 'company' WHERE entity_type = 'contact'")
        )
    if _table_exists(conn, "addresses"):
        conn.execute(
            sa.text("UPDATE addresses SET entity_type = 'company' WHERE entity_type = 'contact'")
        )