#!/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 # 只处理指定素材库 python scripts/recount_asset_counts.py --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()