xiaoxia-saas/scripts/recount_asset_counts.py

#!/usr/bin/env python3
"""素材库计数重算脚本（Issue #1776）。

用法:
    # Dry-run: 输出差异清单，不执行修改
    python scripts/recount_asset_counts.py --dry-run

    # 执行修正
    python scripts/recount_asset_counts.py

    # 只处理指定项目
    python scripts/recount_asset_counts.py --project-id <project_id>

    # 只处理指定素材库
    python scripts/recount_asset_counts.py --library-id <library_id>
"""

import argparse
import sys
from pathlib import Path

# 添加项目根目录到 path
sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy import create_engine, func
from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import AssetLibraryModel, AssetModel


def get_db_session() -> Session:
    """创建数据库 session。"""
    import os

    database_url = os.getenv("DATABASE_URL")
    if not database_url:
        print("ERROR: DATABASE_URL environment variable not set")
        sys.exit(1)
    engine = create_engine(database_url)
    return Session(engine)


def check_discrepancies(session: Session, project_id: str | None = None, library_id: str | None = None) -> list[dict]:
    """检查素材库计数差异。

    返回列表，每项包含:
    - library_id: 素材库 ID
    - library_name: 素材库名称
    - recorded_count: 记录的计数
    - actual_count: 实际计数
    - delta: 差异 (actual - recorded)
    """
    query = session.query(AssetLibraryModel)
    if project_id:
        query = query.filter(AssetLibraryModel.project_id == project_id)
    if library_id:
        query = query.filter(AssetLibraryModel.id == library_id)

    libraries = query.all()
    discrepancies = []

    for lib in libraries:
        # 查询实际计数（排除 deleted）
        actual_count = (
            session.query(func.count(AssetModel.id))
            .filter(
                AssetModel.asset_library_id == lib.id,
                AssetModel.status != "deleted",
            )
            .scalar()
            or 0
        )
        actual_size = (
            session.query(func.coalesce(func.sum(AssetModel.file_size), 0))
            .filter(
                AssetModel.asset_library_id == lib.id,
                AssetModel.status != "deleted",
            )
            .scalar()
            or 0
        )
        recorded_count = int(lib.asset_count or 0)
        recorded_size = int(lib.total_size or 0)

        if actual_count != recorded_count or actual_size != recorded_size:
            discrepancies.append(
                {
                    "library_id": lib.id,
                    "library_name": lib.name,
                    "project_id": lib.project_id,
                    "kind": lib.kind,
                    "recorded_count": recorded_count,
                    "actual_count": actual_count,
                    "count_delta": actual_count - recorded_count,
                    "recorded_size": recorded_size,
                    "actual_size": actual_size,
                    "size_delta": actual_size - recorded_size,
                }
            )

    return discrepancies


def fix_discrepancies(session: Session, discrepancies: list[dict]) -> int:
    """修正素材库计数。返回修正数量。"""
    fixed = 0
    for d in discrepancies:
        session.query(AssetLibraryModel).filter(AssetLibraryModel.id == d["library_id"]).update(
            {
                AssetLibraryModel.asset_count: d["actual_count"],
                AssetLibraryModel.total_size: d["actual_size"],
            }
        )
        fixed += 1
    session.commit()
    return fixed


def main():
    parser = argparse.ArgumentParser(description="素材库计数重算脚本（Issue #1776）")
    parser.add_argument("--dry-run", action="store_true", help="只输出差异清单，不执行修正")
    parser.add_argument("--project-id", type=str, help="只处理指定项目")
    parser.add_argument("--library-id", type=str, help="只处理指定素材库")
    args = parser.parse_args()

    session = get_db_session()

    try:
        discrepancies = check_discrepancies(session, args.project_id, args.library_id)

        if not discrepancies:
            print("✅ 所有素材库计数一致，无需修正")
            return

        # 输出差异清单
        print(f"发现 {len(discrepancies)} 个素材库计数不一致:\n")
        print(f"{'Library ID':<40} {'Name':<20} {'Recorded':<10} {'Actual':<10} {'Delta':<10}")
        print("-" * 90)
        for d in discrepancies:
            print(
                f"{d['library_id']:<40} {d['library_name'][:20]:<20} {d['recorded_count']:<10} {d['actual_count']:<10} {d['count_delta']:+<10}"
            )

        total_delta = sum(d["count_delta"] for d in discrepancies)
        print(f"\n总计差异: {total_delta:+d}")

        if args.dry_run:
            print("\n[DRY-RUN] 未执行修正。移除 --dry-run 参数以执行修正。")
        else:
            print("\n正在执行修正...")
            fixed = fix_discrepancies(session, discrepancies)
            print(f"✅ 已修正 {fixed} 个素材库计数")

    finally:
        session.close()


if __name__ == "__main__":
    main()