geMoldInsight/scripts/migrations/backfill_mold_cavity_confidence.py

#!/usr/bin/env python3
"""
回填 mold_cavity_data 的多方案可信化字段。

回填字段：
- best_scheme_id
- confidence_score
- is_fallback
- fallback_reason
"""

import asyncio
import json
import sys
from pathlib import Path

project_root = Path(__file__).resolve().parents[2]
src_root = project_root / "src"
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(src_root))

from sqlalchemy import select

from database.database import db_manager
from moldinsight.models import MoldCavityData
from storage.rustfs_storage import rustfs_manager
from config.settings import settings
from services.storage_integration_rustfs import StorageIntegrationService
from utils.logger import get_logger

logger = get_logger(__name__)


async def backfill():
    await db_manager.connect()
    await rustfs_manager.connect(
        endpoint=settings.RUSTFS_ENDPOINT,
        access_key=settings.RUSTFS_ACCESS_KEY,
        secret_key=settings.RUSTFS_SECRET_KEY,
        timeout=settings.RUSTFS_TIMEOUT,
    )

    summary = {"total": 0, "updated": 0, "failed": 0}
    try:
        async with db_manager.session() as session:
            result = await session.execute(select(MoldCavityData))
            rows = result.scalars().all()
            summary["total"] = len(rows)

            for row in rows:
                try:
                    cavity_bytes = await rustfs_manager.download_file(
                        file_type="mold_cavities",
                        object_key=row.detailed_object_key,
                    )
                    cavity_json = json.loads(cavity_bytes.decode("utf-8"))
                    payload = StorageIntegrationService._resolve_best_scheme_payload(cavity_json)
                    best_scheme = payload.get("best_scheme") or {}

                    row.best_scheme_id = payload.get("best_scheme_id")
                    row.confidence_score = best_scheme.get("confidence_score")
                    row.is_fallback = best_scheme.get("is_fallback")
                    row.fallback_reason = best_scheme.get("fallback_reason")
                    summary["updated"] += 1
                except Exception as exc:
                    summary["failed"] += 1
                    logger.warning(f"回填失败 id={row.id}: {exc}")

            await session.commit()
    finally:
        await rustfs_manager.close()
        await db_manager.disconnect()

    print("=== Backfill 完成 ===")
    print(f"总记录: {summary['total']}")
    print(f"成功更新: {summary['updated']}")
    print(f"失败: {summary['failed']}")


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