SupplyAgent/alembic/versions/20260731_13_version_demand_belong_rel.py

"""version and explain legacy demand-belong relation edges

Revision ID: 20260731_13
Revises: 20260731_12
Create Date: 2026-07-31
"""
from __future__ import annotations

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op

revision: str = "20260731_13"
down_revision: str | None = "20260731_12"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def _new_columns() -> tuple[sa.Column, ...]:
    return (
        sa.Column(
            "biz_dt",
            sa.String(32),
            nullable=False,
            server_default="legacy",
        ),
        sa.Column(
            "relation_type",
            sa.String(32),
            nullable=False,
            server_default="explicit_token",
        ),
        sa.Column(
            "relation_source",
            sa.String(64),
            nullable=False,
            server_default="legacy",
        ),
        sa.Column(
            "reason",
            sa.Text(),
            nullable=False,
        ),
        sa.Column(
            "confidence",
            sa.Numeric(6, 5),
            nullable=False,
            server_default="1",
        ),
        sa.Column(
            "is_inferred",
            sa.Boolean(),
            nullable=False,
            server_default="1",
        ),
        sa.Column(
            "status",
            sa.String(24),
            nullable=False,
            server_default="active",
        ),
        sa.Column(
            "valid_from_biz_dt",
            sa.String(32),
            nullable=False,
            server_default="legacy",
        ),
        sa.Column("valid_to_biz_dt", sa.String(32), nullable=True),
    )


def upgrade() -> None:
    inspector = sa.inspect(op.get_bind())
    table_names = set(inspector.get_table_names())
    if "demand_belong_pool_rel" not in table_names:
        op.create_table(
            "demand_belong_pool_rel",
            sa.Column("id", sa.BigInteger(), autoincrement=True, nullable=False),
            sa.Column("demand_belong_category_id", sa.BigInteger(), nullable=False),
            sa.Column("multi_demand_pool_di_id", sa.BigInteger(), nullable=False),
            *_new_columns(),
            sa.Column(
                "create_time",
                sa.DateTime(),
                nullable=False,
                server_default=sa.func.now(),
            ),
            sa.Column(
                "update_time",
                sa.DateTime(),
                nullable=False,
                server_default=sa.func.now(),
            ),
            sa.PrimaryKeyConstraint("id"),
            sa.UniqueConstraint(
                "demand_belong_category_id",
                "multi_demand_pool_di_id",
                name="uk_belong_pool",
            ),
        )
        op.create_index(
            "idx_belong_category_id",
            "demand_belong_pool_rel",
            ["demand_belong_category_id"],
        )
        op.create_index(
            "idx_pool_di_id",
            "demand_belong_pool_rel",
            ["multi_demand_pool_di_id"],
        )
    else:
        existing_columns = {
            column["name"]
            for column in inspector.get_columns("demand_belong_pool_rel")
        }
        reason_added = False
        for column in _new_columns():
            if column.name not in existing_columns:
                if column.name == "reason":
                    op.add_column(
                        "demand_belong_pool_rel",
                        sa.Column("reason", sa.Text(), nullable=True),
                    )
                    reason_added = True
                else:
                    op.add_column("demand_belong_pool_rel", column)
        if "multi_demand_pool_di" in table_names:
            op.execute(
                sa.text(
                    "UPDATE demand_belong_pool_rel "
                    "SET biz_dt = COALESCE(("
                    "SELECT pool.biz_dt FROM multi_demand_pool_di AS pool "
                    "WHERE pool.id = demand_belong_pool_rel.multi_demand_pool_di_id"
                    "), 'legacy'), "
                    "valid_from_biz_dt = COALESCE(("
                    "SELECT pool.biz_dt FROM multi_demand_pool_di AS pool "
                    "WHERE pool.id = demand_belong_pool_rel.multi_demand_pool_di_id"
                    "), 'legacy'), "
                    "relation_source = 'legacy_name_match', "
                    "reason = '历史名称匹配关系，由迁移接管；后续按业务日重建', "
                    "is_inferred = 1"
                )
            )
        if reason_added:
            op.execute(
                sa.text(
                    "UPDATE demand_belong_pool_rel "
                    "SET reason = 'legacy relation' "
                    "WHERE reason IS NULL"
                )
            )
            if op.get_bind().dialect.name == "sqlite":
                with op.batch_alter_table("demand_belong_pool_rel") as batch_op:
                    batch_op.alter_column(
                        "reason",
                        existing_type=sa.Text(),
                        nullable=False,
                    )
            else:
                op.alter_column(
                    "demand_belong_pool_rel",
                    "reason",
                    existing_type=sa.Text(),
                    nullable=False,
                )
    index_names = {
        index["name"]
        for index in sa.inspect(op.get_bind()).get_indexes(
            "demand_belong_pool_rel"
        )
    }
    if "idx_belong_pool_biz_status" not in index_names:
        op.create_index(
            "idx_belong_pool_biz_status",
            "demand_belong_pool_rel",
            ["biz_dt", "status"],
        )


def downgrade() -> None:
    op.drop_index(
        "idx_belong_pool_biz_status",
        table_name="demand_belong_pool_rel",
    )
    for column in (
        "valid_to_biz_dt",
        "valid_from_biz_dt",
        "status",
        "is_inferred",
        "confidence",
        "reason",
        "relation_source",
        "relation_type",
        "biz_dt",
    ):
        op.drop_column("demand_belong_pool_rel", column)