video-shuoshan/backend/apps/generation/management/commands/sync_volcano_assets.py

"""
从火山 default 项目同步素材组+素材到本地数据库。

用法:
    python manage.py sync_volcano_assets --team-id 1 [--dry-run]
    python manage.py sync_volcano_assets --team-id 1 --group-ids group-xxx,group-yyy
"""
import logging
from django.core.management.base import BaseCommand
from apps.accounts.models import Team
from apps.generation.models import AssetGroup, Asset
from utils.assets_client import _do_request

logger = logging.getLogger(__name__)

# 火山素材所在的项目名
SOURCE_PROJECT = 'default'


def fetch_groups_from_volcano(group_ids=None, page_size=50):
    """从火山 default 项目拉取资源组列表。"""
    body = {
        'Filter': {'GroupType': 'AIGC'},
        'PageNumber': 1,
        'PageSize': page_size,
        'ProjectName': SOURCE_PROJECT,
    }
    if group_ids:
        body['Filter']['GroupIds'] = group_ids
    result = _do_request('ListAssetGroups', body)
    return result.get('Items', [])


def fetch_assets_from_volcano(group_id):
    """从火山 default 项目拉取某组下的所有素材。"""
    body = {
        'Filter': {'GroupType': 'AIGC', 'GroupIds': [group_id]},
        'PageNumber': 1,
        'PageSize': 100,
        'ProjectName': SOURCE_PROJECT,
    }
    result = _do_request('ListAssets', body)
    return result.get('Items', [])


class Command(BaseCommand):
    help = '从火山 default 项目同步素材组和素材到本地数据库'

    def add_arguments(self, parser):
        group = parser.add_mutually_exclusive_group(required=True)
        group.add_argument('--team-id', type=int, help='绑定到的团队 ID')
        group.add_argument('--team-name', type=str, help='绑定到的团队名称')
        parser.add_argument('--group-ids', type=str, default='', help='指定要同步的组 ID（逗号分隔），留空同步全部')
        parser.add_argument('--dry-run', action='store_true', help='只打印不写入数据库')

    def handle(self, *args, **options):
        dry_run = options['dry_run']
        group_ids_str = options['group_ids'].strip()

        if options['team_id']:
            try:
                team = Team.objects.get(pk=options['team_id'])
            except Team.DoesNotExist:
                self.stderr.write(f'团队 ID {options["team_id"]} 不存在')
                return
        else:
            try:
                team = Team.objects.get(name=options['team_name'])
            except Team.DoesNotExist:
                self.stderr.write(f'团队「{options["team_name"]}」不存在')
                return

        self.stdout.write(f'目标团队: {team.name} (ID={team_id})')
        if dry_run:
            self.stdout.write('** DRY RUN — 不会写入数据库 **')

        # 拉取火山资源组
        group_ids = [g.strip() for g in group_ids_str.split(',') if g.strip()] if group_ids_str else None
        volcano_groups = fetch_groups_from_volcano(group_ids=group_ids)
        self.stdout.write(f'从火山 {SOURCE_PROJECT} 项目拉取到 {len(volcano_groups)} 个资源组')

        created_groups = 0
        created_assets = 0
        skipped_groups = 0
        skipped_assets = 0

        for vg in volcano_groups:
            remote_gid = vg.get('Id', '')
            group_name = vg.get('Name', '') or remote_gid

            # 检查本地是否已存在
            existing = AssetGroup.objects.filter(remote_group_id=remote_gid, team=team).first()
            if existing:
                self.stdout.write(f'  [跳过] 组 {group_name} ({remote_gid}) — 本地已存在 (pk={existing.pk})')
                skipped_groups += 1
                local_group = existing
            else:
                self.stdout.write(f'  [新增] 组 {group_name} ({remote_gid})')
                if not dry_run:
                    local_group = AssetGroup.objects.create(
                        team=team,
                        remote_group_id=remote_gid,
                        name=group_name,
                        description=vg.get('Description', ''),
                    )
                else:
                    local_group = None
                created_groups += 1

            # 拉取该组下的素材
            volcano_assets = fetch_assets_from_volcano(remote_gid)
            self.stdout.write(f'    └ {len(volcano_assets)} 个素材')

            for va in volcano_assets:
                remote_aid = va.get('Id', '')
                asset_name = va.get('Name', '') or remote_aid
                asset_type = va.get('AssetType', 'Image')
                status_map = {'Active': 'active', 'Processing': 'processing', 'Failed': 'failed'}
                local_status = status_map.get(va.get('Status', ''), 'processing')
                preview_url = va.get('PreviewUrl', '') or ''

                # 检查本地是否已存在
                if local_group and Asset.objects.filter(remote_asset_id=remote_aid, group=local_group).exists():
                    self.stdout.write(f'      [跳过] 素材 {asset_name} ({remote_aid})')
                    skipped_assets += 1
                    continue

                self.stdout.write(f'      [新增] 素材 {asset_name} | {asset_type} | {local_status} | {remote_aid}')
                if not dry_run and local_group:
                    Asset.objects.create(
                        group=local_group,
                        remote_asset_id=remote_aid,
                        name=asset_name,
                        url=preview_url,
                        asset_type=asset_type,
                        status=local_status,
                    )
                created_assets += 1

        # 更新组缩略图（取第一张图片素材的 URL）
        if not dry_run:
            for ag in AssetGroup.objects.filter(team=team, thumbnail_url=''):
                first_img = ag.assets.filter(asset_type='Image', status='active').first()
                if first_img and first_img.url:
                    ag.thumbnail_url = first_img.url
                    ag.save(update_fields=['thumbnail_url'])
                    self.stdout.write(f'  [更新缩略图] {ag.name}')

        self.stdout.write('')
        self.stdout.write(f'完成! 新增 {created_groups} 组 + {created_assets} 素材，跳过 {skipped_groups} 组 + {skipped_assets} 素材')