dpb/backend/scripts/schema_align.py

"""波次5.1 数据治理：bre_* 表审计列统一到新契约。

新契约（与 core/base_model.py ModelMixin + UserMixin 一致）：
    id / uuid(unique) / is_deleted(bool) / created_time / updated_time / deleted_time /
    created_id / updated_id / deleted_id

处理三类历史表（全部幂等，可重复执行）：
- 旧式表（is_deleted=int，遗留 created_by/updated_by 列，缺 uuid/deleted_time）：
  补 uuid（回填 gen_random_uuid + 唯一索引）、补 deleted_time、
  is_deleted int→boolean、删除 created_by/updated_by（先并入 created_id/updated_id）。
- 缺审计列表（bre_prediction_value / bre_audit_log / bre_statistics_job）：
  补缺列（uuid/is_deleted/deleted_time/updated_time）。
- 新式表残留死列（created_by/updated_by）：直接删除。

说明：
- 旧式表既有 created_id（新）又有 created_by（旧），合并时以 created_id 优先。
- 统计 4 表（prediction/prediction_value/combining_ability/statistics_job）已接 ORM，
  补 created_deleted 索引与 ORM __table_args__ 对齐。
"""
import asyncio
import os
import sys

BACKEND_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, BACKEND_DIR)

os.environ.setdefault("ENVIRONMENT", "dev")

# 旧式表：is_deleted=int + created_by/updated_by 遗留，缺 uuid/deleted_time
GROUP_A = [
    "bre_marker", "bre_genotype_call", "bre_genotype_sample", "bre_genotyping_dataset",
    "bre_environment_condition", "bre_field_operation", "bre_group_member", "bre_observation",
    "bre_planting_treatment", "bre_prediction", "bre_propagation", "bre_report", "bre_seed_lot",
    "bre_combining_ability",
]
# 特殊：bre_statistics_job 缺 uuid/deleted_time/updated_time，且有 created_by 遗留
GROUP_C = ["bre_statistics_job"]
# 缺审计列：uuid/is_deleted/deleted_time/updated_time 全缺（无 created_by 遗留）
GROUP_B = ["bre_prediction_value", "bre_audit_log"]
# 新式表残留 created_by/updated_by 死列
GROUP_D = [
    "bre_group", "bre_pedigree", "bre_plan", "bre_selection_rule", "bre_trait",
    "bre_treatment", "bre_trial", "bre_trial_study", "bre_trial_study_entry",
]
# 统计 4 表补 created_deleted 索引（与 ORM __table_args__ 对齐）
INDEX_TABLES = ["bre_prediction", "bre_prediction_value", "bre_combining_ability", "bre_statistics_job"]


async def column_exists(db, table: str, col: str) -> bool:
    from sqlalchemy import text

    r = await db.execute(text(
        "SELECT 1 FROM information_schema.columns WHERE table_name=:t AND column_name=:c"
    ), {"t": table, "c": col})
    return r.scalar() is not None


async def index_exists(db, table: str, idx: str) -> bool:
    from sqlalchemy import text

    r = await db.execute(text(
        "SELECT 1 FROM pg_indexes WHERE schemaname='public' AND tablename=:t AND indexname=:i"
    ), {"t": table, "i": idx})
    return r.scalar() is not None


async def column_type(db, table: str, col: str) -> str:
    from sqlalchemy import text

    r = await db.execute(text(
        "SELECT data_type FROM information_schema.columns WHERE table_name=:t AND column_name=:c"
    ), {"t": table, "c": col})
    return (r.scalar() or "").lower()


async def ensure_uuid(db, table: str) -> None:
    from sqlalchemy import text

    if not await column_exists(db, table, "uuid"):
        await db.execute(text(f"ALTER TABLE {table} ADD COLUMN uuid VARCHAR(64)"))
    await db.execute(text(
        f"UPDATE {table} SET uuid = gen_random_uuid()::varchar WHERE uuid IS NULL OR btrim(uuid) = ''"
    ))
    await db.execute(text(f"ALTER TABLE {table} ALTER COLUMN uuid SET NOT NULL"))
    idx = f"ix_{table}_uuid"
    if not await index_exists(db, table, idx):
        await db.execute(text(f"CREATE UNIQUE INDEX {idx} ON {table} (uuid)"))


async def ensure_column(db, table: str, col: str, ddl: str) -> None:
    from sqlalchemy import text

    if not await column_exists(db, table, col):
        await db.execute(text(ddl))


async def drop_legacy_audit(db, table: str) -> None:
    """合并旧列数据到新列后删除 created_by/updated_by 死列。"""
    from sqlalchemy import text

    if await column_exists(db, table, "created_by"):
        await db.execute(text(
            f"UPDATE {table} SET created_id = created_by "
            f"WHERE created_id IS NULL AND created_by IS NOT NULL"
        ))
        await db.execute(text(f"ALTER TABLE {table} DROP COLUMN created_by"))
    if await column_exists(db, table, "updated_by"):
        await db.execute(text(
            f"UPDATE {table} SET updated_id = updated_by "
            f"WHERE updated_id IS NULL AND updated_by IS NOT NULL"
        ))
        await db.execute(text(f"ALTER TABLE {table} DROP COLUMN updated_by"))


async def convert_is_deleted(db, table: str) -> None:
    from sqlalchemy import text

    t = await column_type(db, table, "is_deleted")
    if t != "integer":
        return
    bad = (await db.execute(text(
        f"SELECT COUNT(*) FROM {table} WHERE is_deleted IS NULL OR is_deleted NOT IN (0,1)"
    ))).scalar() or 0
    if bad:
        raise RuntimeError(f"{table}.is_deleted 含非 0/1 值 {bad} 行，拒绝转换")
    # 先去掉整数默认值（0），否则类型转换会因默认值无法转换而失败
    await db.execute(text(f"ALTER TABLE {table} ALTER COLUMN is_deleted DROP DEFAULT"))
    await db.execute(text(
        f"ALTER TABLE {table} ALTER COLUMN is_deleted TYPE BOOLEAN USING (is_deleted::int::boolean)"
    ))
    await db.execute(text(f"ALTER TABLE {table} ALTER COLUMN is_deleted SET DEFAULT false"))


async def ensure_created_deleted_index(db, table: str) -> None:
    from sqlalchemy import text

    idx = f"ix_{table}_created_deleted"
    if not await index_exists(db, table, idx):
        await db.execute(text(
            f"CREATE INDEX {idx} ON {table} (created_time, is_deleted)"
        ))


async def main() -> None:
    from sqlalchemy import text

    from app.core.database import async_db_session

    done: list[str] = []
    async with async_db_session() as db:
        for t in GROUP_A:
            await ensure_uuid(db, t)
            await ensure_column(db, t, "deleted_time", f"ALTER TABLE {t} ADD COLUMN deleted_time TIMESTAMPTZ")
            await convert_is_deleted(db, t)
            await drop_legacy_audit(db, t)
            await db.commit()
            done.append(f"{t} (old-style)")

        for t in GROUP_C:
            await ensure_uuid(db, t)
            await ensure_column(db, t, "deleted_time", f"ALTER TABLE {t} ADD COLUMN deleted_time TIMESTAMPTZ")
            await convert_is_deleted(db, t)
            await drop_legacy_audit(db, t)
            if not await column_exists(db, t, "updated_time"):
                await db.execute(text(
                    f"ALTER TABLE {t} ADD COLUMN updated_time TIMESTAMPTZ DEFAULT now() NOT NULL"
                ))
            await db.commit()
            done.append(f"{t} (old-style, 补 updated_time)")

        for t in GROUP_B:
            await ensure_uuid(db, t)
            await ensure_column(db, t, "is_deleted", f"ALTER TABLE {t} ADD COLUMN is_deleted BOOLEAN DEFAULT false NOT NULL")
            await ensure_column(db, t, "deleted_time", f"ALTER TABLE {t} ADD COLUMN deleted_time TIMESTAMPTZ")
            if not await column_exists(db, t, "updated_time"):
                await db.execute(text(
                    f"ALTER TABLE {t} ADD COLUMN updated_time TIMESTAMPTZ DEFAULT now() NOT NULL"
                ))
            await db.commit()
            done.append(f"{t} (缺审计列)")

        for t in GROUP_D:
            await drop_legacy_audit(db, t)
            await db.commit()
            done.append(f"{t} (删死列)")

        for t in INDEX_TABLES:
            await ensure_created_deleted_index(db, t)
            await db.commit()
            done.append(f"{t} (created_deleted 索引)")

    print("完成 " + str(len(done)) + " 项迁移：")
    for d in done:
        print("  - " + d)


if __name__ == "__main__":
    asyncio.run(main())