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