WonderQ-Project/WonderQ-Admin/alembic/versions/0024_vehicle_demand.py

"""Add vehicle demand fields and service configuration."""

from alembic import context, op
import sqlalchemy as sa
from sqlalchemy import inspect
from sqlalchemy.dialects import postgresql

from app.migration_compat import (
    baseline_column_names,
    baseline_foreign_keys,
    baseline_index_names,
    baseline_table_names,
)


revision = "0024_vehicle_demand"
down_revision = "0023_detail_concierge_advisor"
branch_labels = None
depends_on = None


LEAD_COLUMNS = {
    "leadType",
    "contactName",
    "customerId",
    "vehicleDemand",
}

LEAD_INDEXES = {
    "ix_Lead_leadType": ["leadType"],
    "ix_Lead_customerId": ["customerId"],
}

VEHICLE_SERVICE_INDEXES = {
    "ix_VehicleServiceConfig_isActive": ["isActive"],
}


def missing_lead_columns(existing_columns: set[str]) -> set[str]:
    return LEAD_COLUMNS.difference(existing_columns)


def missing_named_indexes(existing_indexes: set[str], expected_indexes: dict[str, list[str]]) -> set[str]:
    return set(expected_indexes).difference(existing_indexes)


def should_create_customer_fk(foreign_keys: list[dict]) -> bool:
    return not any(
        foreign_key.get("referred_table") == "Customer"
        and foreign_key.get("referred_columns") == ["id"]
        and foreign_key.get("constrained_columns") == ["customerId"]
        for foreign_key in foreign_keys
    )


def is_offline_mode() -> bool:
    try:
        return context.is_offline_mode()
    except (NameError, AttributeError):
        return False


def upgrade() -> None:
    if is_offline_mode():
        existing_lead_columns = baseline_column_names("Lead")
        existing_lead_indexes = baseline_index_names("Lead")
        existing_lead_foreign_keys = baseline_foreign_keys("Lead")
        vehicle_service_table_exists = "VehicleServiceConfig" in baseline_table_names()
        existing_vehicle_service_indexes = set()
    else:
        inspector = inspect(op.get_bind())
        existing_lead_columns = {
            column["name"]
            for column in inspector.get_columns("Lead")
        }
        existing_lead_indexes = {
            index["name"]
            for index in inspector.get_indexes("Lead")
            if index.get("name")
        }
        existing_lead_foreign_keys = inspector.get_foreign_keys("Lead")
        vehicle_service_table_exists = "VehicleServiceConfig" in set(inspector.get_table_names())
        existing_vehicle_service_indexes = {
            index["name"]
            for index in inspector.get_indexes("VehicleServiceConfig")
            if index.get("name")
        } if vehicle_service_table_exists else set()

    added_lead_type = "leadType" not in existing_lead_columns
    lead_columns = {
        "leadType": sa.Column("leadType", sa.String(), nullable=False, server_default="general"),
        "contactName": sa.Column("contactName", sa.String(), nullable=True),
        "customerId": sa.Column("customerId", sa.String(), nullable=True),
        "vehicleDemand": sa.Column("vehicleDemand", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
    }
    for column_name in ("leadType", "contactName", "customerId", "vehicleDemand"):
        if column_name in missing_lead_columns(existing_lead_columns):
            op.add_column("Lead", lead_columns[column_name])

    for index_name in missing_named_indexes(existing_lead_indexes, LEAD_INDEXES):
        op.create_index(index_name, "Lead", LEAD_INDEXES[index_name])

    if should_create_customer_fk(existing_lead_foreign_keys):
        op.create_foreign_key(
            "fk_Lead_customerId_Customer",
            "Lead",
            "Customer",
            ["customerId"],
            ["id"],
            ondelete="SET NULL",
        )

    if added_lead_type:
        op.alter_column("Lead", "leadType", server_default=None)

    if not vehicle_service_table_exists:
        op.create_table(
            "VehicleServiceConfig",
            sa.Column("id", sa.String(), nullable=False),
            sa.Column("introTitle", sa.String(), nullable=False),
            sa.Column("intro", sa.Text(), nullable=False),
            sa.Column("serviceSections", postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default=sa.text("'[]'::jsonb")),
            sa.Column("advantages", postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default=sa.text("'[]'::jsonb")),
            sa.Column("processSteps", postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default=sa.text("'[]'::jsonb")),
            sa.Column("isActive", sa.Boolean(), nullable=False, server_default=sa.text("true")),
            sa.Column("createdAt", sa.DateTime(), nullable=False),
            sa.Column("updatedAt", sa.DateTime(), nullable=False),
            sa.PrimaryKeyConstraint("id"),
        )
        existing_vehicle_service_indexes = set()
        for column_name in ("serviceSections", "advantages", "processSteps", "isActive"):
            if column_name != "isActive":
                op.alter_column("VehicleServiceConfig", column_name, server_default=None)
        op.alter_column("VehicleServiceConfig", "isActive", server_default=None)

    for index_name in missing_named_indexes(existing_vehicle_service_indexes, VEHICLE_SERVICE_INDEXES):
        op.create_index(index_name, "VehicleServiceConfig", VEHICLE_SERVICE_INDEXES[index_name])


def downgrade() -> None:
    op.drop_index("ix_VehicleServiceConfig_isActive", table_name="VehicleServiceConfig")
    op.drop_table("VehicleServiceConfig")
    op.drop_constraint("fk_Lead_customerId_Customer", "Lead", type_="foreignkey")
    op.drop_index("ix_Lead_customerId", table_name="Lead")
    op.drop_index("ix_Lead_leadType", table_name="Lead")
    op.drop_column("Lead", "vehicleDemand")
    op.drop_column("Lead", "customerId")
    op.drop_column("Lead", "contactName")
    op.drop_column("Lead", "leadType")