369 lines
15 KiB
Python
369 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict, dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.enums.upload_resource import (
|
|
UploadResourceBindStatusEnum,
|
|
UploadResourceCreatedByEnum,
|
|
UploadResourceDeletePolicyEnum,
|
|
UploadResourceEventEnum,
|
|
UploadResourceModuleEnum,
|
|
UploadResourceSourceModelEnum,
|
|
)
|
|
from app.models.module_generation_project import ModuleGenerationProject
|
|
from app.models.module_generation_step import ModuleGenerationStep
|
|
from app.models.shot_replicate_segment import ShotReplicateSegment
|
|
from app.models.shot_replicate_task_set import ShotReplicateTaskSet
|
|
from app.models.upload_resource import UploadResource
|
|
from app.models.user import User
|
|
from app.services.upload_resource.accounting_service import rebuild_user_resource_stats
|
|
from app.services.upload_resource.bind_service import bind_upload_resources
|
|
from app.services.upload_resource.core_service import record_external_upload_resource
|
|
from app.services.upload_resource.log_service import log_upload_resource_event
|
|
from app.services.upload_resource.file_cleanup_service import cleanup_pending_upload_resource_files
|
|
from app.services.upload_resource.path_resolver import ParsedUploadPath, iter_files, parse_upload_path, upload_url_to_storage_path
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class BackfillOptions:
|
|
root: str = "storage/uploads"
|
|
batch_size: int = 500
|
|
dry_run: bool = True
|
|
include_legacy: bool = False
|
|
rebind_modules: bool = False
|
|
rebuild_stats: bool = False
|
|
only_user_id: str | None = None
|
|
only_module: str | None = None
|
|
cleanup_pending_files: bool = False
|
|
cleanup_limit: int = 500
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class BackfillResult:
|
|
scanned: int = 0
|
|
matched: int = 0
|
|
inserted: int = 0
|
|
updated: int = 0
|
|
existed: int = 0
|
|
skipped: int = 0
|
|
rebind_bound: int = 0
|
|
rebuild_users: int = 0
|
|
cleanup_files: int = 0
|
|
cleanup_failed: int = 0
|
|
skip_reasons: dict[str, int] = field(default_factory=dict)
|
|
|
|
def skip(self, reason: str) -> None:
|
|
self.skipped += 1
|
|
self.skip_reasons[reason] = self.skip_reasons.get(reason, 0) + 1
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"scanned": self.scanned,
|
|
"matched": self.matched,
|
|
"inserted": self.inserted,
|
|
"updated": self.updated,
|
|
"existed": self.existed,
|
|
"skipped": self.skipped,
|
|
"rebind_bound": self.rebind_bound,
|
|
"rebuild_users": self.rebuild_users,
|
|
"cleanup_files": self.cleanup_files,
|
|
"cleanup_failed": self.cleanup_failed,
|
|
"skip_reasons": self.skip_reasons,
|
|
}
|
|
|
|
|
|
def _module_allowed(parsed: ParsedUploadPath, only_module: str | None) -> bool:
|
|
if not only_module:
|
|
return True
|
|
if only_module == UploadResourceModuleEnum.COMMON.value:
|
|
return parsed.module == UploadResourceModuleEnum.COMMON.value
|
|
if only_module == UploadResourceModuleEnum.SHOT_REPLICATE.value:
|
|
return parsed.module == UploadResourceModuleEnum.SHOT_REPLICATE.value
|
|
return parsed.module == only_module
|
|
|
|
|
|
def _collect_urls(value: Any) -> list[str]:
|
|
urls: list[str] = []
|
|
if value is None:
|
|
return urls
|
|
if isinstance(value, str):
|
|
if value.startswith("/uploads/"):
|
|
urls.append(value)
|
|
return urls
|
|
if isinstance(value, dict):
|
|
for v in value.values():
|
|
urls.extend(_collect_urls(v))
|
|
return urls
|
|
if isinstance(value, list):
|
|
for item in value:
|
|
urls.extend(_collect_urls(item))
|
|
return urls
|
|
|
|
|
|
async def _user_exists_map(db: AsyncSession, user_ids: Iterable[str]) -> set[str]:
|
|
ids = [v for v in dict.fromkeys(user_ids) if v]
|
|
if not ids:
|
|
return set()
|
|
result = await db.execute(select(User.id).where(User.id.in_(ids)))
|
|
return set(result.scalars().all())
|
|
|
|
|
|
async def _segment_map(db: AsyncSession, segment_ids: Iterable[str]) -> dict[str, ShotReplicateSegment]:
|
|
ids = [v for v in dict.fromkeys(segment_ids) if v]
|
|
if not ids:
|
|
return {}
|
|
result = await db.execute(select(ShotReplicateSegment).where(ShotReplicateSegment.id.in_(ids)))
|
|
return {item.id: item for item in result.scalars().all()}
|
|
|
|
|
|
async def _process_batch(db: AsyncSession, parsed_items: list[ParsedUploadPath], result: BackfillResult, options: BackfillOptions) -> None:
|
|
if not parsed_items:
|
|
return
|
|
|
|
segment_ids = [p.source_id for p in parsed_items if p.resource_type == "shot_segment" and p.source_id]
|
|
segments = await _segment_map(db, segment_ids)
|
|
|
|
for parsed in parsed_items:
|
|
if parsed.resource_type == "shot_segment":
|
|
segment = segments.get(parsed.source_id or "")
|
|
if not segment:
|
|
result.skip("shot_segment_not_found")
|
|
continue
|
|
parsed.user_id = segment.user_id
|
|
parsed.created_at = segment.created_at or parsed.created_at
|
|
parsed.source_model = UploadResourceSourceModelEnum.SHOT_REPLICATE_SEGMENT.value
|
|
parsed.source_id = segment.id
|
|
|
|
users = await _user_exists_map(db, [p.user_id for p in parsed_items if p.user_id])
|
|
existing_result = await db.execute(select(UploadResource.storage_path).where(UploadResource.storage_path.in_([p.storage_path for p in parsed_items])))
|
|
existing_paths = set(existing_result.scalars().all())
|
|
|
|
for parsed in parsed_items:
|
|
if not parsed.user_id:
|
|
result.skip("missing_user_id")
|
|
continue
|
|
if options.only_user_id and parsed.user_id != options.only_user_id:
|
|
result.skip("user_filtered")
|
|
continue
|
|
if parsed.user_id not in users:
|
|
result.skip("user_not_found")
|
|
continue
|
|
if parsed.storage_path in existing_paths:
|
|
result.existed += 1
|
|
if not options.dry_run:
|
|
await record_external_upload_resource(
|
|
db,
|
|
user_id=parsed.user_id,
|
|
module=parsed.module,
|
|
resource_type=parsed.resource_type,
|
|
resource_url=parsed.resource_url,
|
|
storage_path=parsed.storage_path,
|
|
file_size_bytes=parsed.file_size_bytes,
|
|
file_name=parsed.file_name,
|
|
source_model=parsed.source_model,
|
|
source_id=parsed.source_id,
|
|
bind_status=UploadResourceBindStatusEnum.BOUND.value if parsed.source_id else UploadResourceBindStatusEnum.PENDING.value,
|
|
delete_policy=UploadResourceDeletePolicyEnum.MODULE_ONLY.value if parsed.source_id else UploadResourceDeletePolicyEnum.USER_DELETABLE.value,
|
|
created_by=UploadResourceCreatedByEnum.BACKFILL.value,
|
|
created_at=parsed.created_at,
|
|
metadata={"backfill": True},
|
|
)
|
|
result.updated += 1
|
|
continue
|
|
result.matched += 1
|
|
if options.dry_run:
|
|
continue
|
|
await record_external_upload_resource(
|
|
db,
|
|
user_id=parsed.user_id,
|
|
module=parsed.module,
|
|
resource_type=parsed.resource_type,
|
|
resource_url=parsed.resource_url,
|
|
storage_path=parsed.storage_path,
|
|
file_size_bytes=parsed.file_size_bytes,
|
|
file_name=parsed.file_name,
|
|
source_model=parsed.source_model,
|
|
source_id=parsed.source_id,
|
|
bind_status=UploadResourceBindStatusEnum.BOUND.value if parsed.source_id else UploadResourceBindStatusEnum.PENDING.value,
|
|
delete_policy=UploadResourceDeletePolicyEnum.MODULE_ONLY.value if parsed.source_id else UploadResourceDeletePolicyEnum.USER_DELETABLE.value,
|
|
created_by=UploadResourceCreatedByEnum.BACKFILL.value,
|
|
created_at=parsed.created_at,
|
|
metadata={"backfill": True},
|
|
)
|
|
result.inserted += 1
|
|
|
|
|
|
def _batched(items: Iterable[Path], batch_size: int):
|
|
batch: list[Path] = []
|
|
for item in items:
|
|
batch.append(item)
|
|
if len(batch) >= batch_size:
|
|
yield batch
|
|
batch = []
|
|
if batch:
|
|
yield batch
|
|
|
|
|
|
async def _rebind_hot_opening(db: AsyncSession, *, only_user_id: str | None = None) -> int:
|
|
stmt = select(ModuleGenerationProject.id, ModuleGenerationProject.user_id).where(
|
|
ModuleGenerationProject.module == UploadResourceModuleEnum.HOT_OPENING_REPLICATE.value,
|
|
ModuleGenerationProject.deleted_at.is_(None),
|
|
)
|
|
if only_user_id:
|
|
stmt = stmt.where(ModuleGenerationProject.user_id == only_user_id)
|
|
projects = (await db.execute(stmt)).all()
|
|
if not projects:
|
|
return 0
|
|
project_user = {pid: uid for pid, uid in projects}
|
|
steps_result = await db.execute(
|
|
select(ModuleGenerationStep.project_id, ModuleGenerationStep.input_json, ModuleGenerationStep.output_json).where(
|
|
ModuleGenerationStep.project_id.in_(list(project_user.keys())),
|
|
ModuleGenerationStep.deleted_at.is_(None),
|
|
)
|
|
)
|
|
grouped: dict[tuple[str, str], list[str]] = {}
|
|
for project_id, input_json, output_json in steps_result.all():
|
|
urls = _collect_urls(input_json) + _collect_urls(output_json)
|
|
key = (project_user[project_id], project_id)
|
|
grouped.setdefault(key, []).extend(urls)
|
|
|
|
bound = 0
|
|
for (user_id, project_id), urls in grouped.items():
|
|
stats = await bind_upload_resources(
|
|
db,
|
|
user_id=user_id,
|
|
module=UploadResourceModuleEnum.HOT_OPENING_REPLICATE.value,
|
|
source_model=UploadResourceSourceModelEnum.MODULE_GENERATION_PROJECT.value,
|
|
source_id=project_id,
|
|
urls=urls,
|
|
allow_common_migrate=True,
|
|
)
|
|
bound += stats.get("bound", 0)
|
|
return bound
|
|
|
|
|
|
async def _rebind_shot_replicate(db: AsyncSession, *, only_user_id: str | None = None) -> int:
|
|
bound = 0
|
|
task_stmt = select(ShotReplicateTaskSet.id, ShotReplicateTaskSet.user_id, ShotReplicateTaskSet.video_url).where(ShotReplicateTaskSet.deleted_at.is_(None))
|
|
if only_user_id:
|
|
task_stmt = task_stmt.where(ShotReplicateTaskSet.user_id == only_user_id)
|
|
for task_set_id, user_id, video_url in (await db.execute(task_stmt)).all():
|
|
stats = await bind_upload_resources(
|
|
db,
|
|
user_id=user_id,
|
|
module=UploadResourceModuleEnum.SHOT_REPLICATE.value,
|
|
source_model=UploadResourceSourceModelEnum.SHOT_REPLICATE_TASK_SET.value,
|
|
source_id=task_set_id,
|
|
urls=[video_url],
|
|
allow_common_migrate=True,
|
|
)
|
|
bound += stats.get("bound", 0)
|
|
|
|
project_stmt = select(ModuleGenerationProject.id, ModuleGenerationProject.user_id).where(
|
|
ModuleGenerationProject.module == UploadResourceModuleEnum.SHOT_REPLICATE.value,
|
|
ModuleGenerationProject.deleted_at.is_(None),
|
|
)
|
|
if only_user_id:
|
|
project_stmt = project_stmt.where(ModuleGenerationProject.user_id == only_user_id)
|
|
projects = (await db.execute(project_stmt)).all()
|
|
if projects:
|
|
project_user = {pid: uid for pid, uid in projects}
|
|
steps_result = await db.execute(
|
|
select(ModuleGenerationStep.project_id, ModuleGenerationStep.input_json, ModuleGenerationStep.output_json).where(
|
|
ModuleGenerationStep.project_id.in_(list(project_user.keys())),
|
|
ModuleGenerationStep.deleted_at.is_(None),
|
|
)
|
|
)
|
|
grouped: dict[tuple[str, str], list[str]] = {}
|
|
for project_id, input_json, output_json in steps_result.all():
|
|
key = (project_user[project_id], project_id)
|
|
grouped.setdefault(key, []).extend(_collect_urls(input_json) + _collect_urls(output_json))
|
|
for (user_id, project_id), urls in grouped.items():
|
|
stats = await bind_upload_resources(
|
|
db,
|
|
user_id=user_id,
|
|
module=UploadResourceModuleEnum.SHOT_REPLICATE.value,
|
|
source_model=UploadResourceSourceModelEnum.MODULE_GENERATION_PROJECT.value,
|
|
source_id=project_id,
|
|
urls=urls,
|
|
allow_common_migrate=True,
|
|
)
|
|
bound += stats.get("bound", 0)
|
|
return bound
|
|
|
|
|
|
async def rebind_module_upload_resources(db: AsyncSession, *, only_module: str | None = None, only_user_id: str | None = None) -> int:
|
|
total = 0
|
|
if only_module in (None, UploadResourceModuleEnum.HOT_OPENING_REPLICATE.value):
|
|
total += await _rebind_hot_opening(db, only_user_id=only_user_id)
|
|
if only_module in (None, UploadResourceModuleEnum.SHOT_REPLICATE.value):
|
|
total += await _rebind_shot_replicate(db, only_user_id=only_user_id)
|
|
return total
|
|
|
|
|
|
async def run_upload_resource_backfill(db: AsyncSession, options: BackfillOptions) -> BackfillResult:
|
|
result = BackfillResult()
|
|
log_upload_resource_event(
|
|
event_type=UploadResourceEventEnum.BACKFILL_START.value,
|
|
detail={"options": asdict(options)},
|
|
)
|
|
|
|
if options.cleanup_pending_files:
|
|
if not options.dry_run:
|
|
cleanup_stats = await cleanup_pending_upload_resource_files(db, limit=options.cleanup_limit)
|
|
result.cleanup_files = int(cleanup_stats.get("deleted", 0) or 0) + int(cleanup_stats.get("missing", 0) or 0)
|
|
result.cleanup_failed = int(cleanup_stats.get("failed", 0) or 0)
|
|
log_upload_resource_event(
|
|
event_type=UploadResourceEventEnum.BACKFILL_FINISHED.value,
|
|
detail=result.to_dict(),
|
|
)
|
|
return result
|
|
|
|
for paths in _batched(iter_files(options.root), options.batch_size):
|
|
parsed_items: list[ParsedUploadPath] = []
|
|
for path in paths:
|
|
result.scanned += 1
|
|
parsed = parse_upload_path(path, include_legacy=options.include_legacy)
|
|
if not parsed:
|
|
result.skip("unparsed")
|
|
continue
|
|
if parsed.skip_reason:
|
|
result.skip(parsed.skip_reason)
|
|
continue
|
|
if not _module_allowed(parsed, options.only_module):
|
|
result.skip("module_filtered")
|
|
continue
|
|
parsed_items.append(parsed)
|
|
await _process_batch(db, parsed_items, result, options)
|
|
if not options.dry_run:
|
|
await db.commit()
|
|
|
|
if options.rebind_modules:
|
|
if options.dry_run:
|
|
# rebind dry-run 不真实改库，只统计为 0，避免复杂模拟误导。
|
|
result.rebind_bound = 0
|
|
else:
|
|
result.rebind_bound = await rebind_module_upload_resources(
|
|
db,
|
|
only_module=options.only_module if options.only_module in (UploadResourceModuleEnum.HOT_OPENING_REPLICATE.value, UploadResourceModuleEnum.SHOT_REPLICATE.value) else None,
|
|
only_user_id=options.only_user_id,
|
|
)
|
|
await db.commit()
|
|
|
|
if options.rebuild_stats and not options.dry_run:
|
|
stats = await rebuild_user_resource_stats(db, user_ids=[options.only_user_id] if options.only_user_id else None)
|
|
result.rebuild_users = stats.get("users", 0)
|
|
await db.commit()
|
|
|
|
log_upload_resource_event(
|
|
event_type=UploadResourceEventEnum.BACKFILL_FINISHED.value,
|
|
detail=result.to_dict(),
|
|
)
|
|
return result
|