异步上传素材

2026-06-22 13:26:43 +08:00
parent f62dd8bf67
commit 9d79dba2c1
4 changed files with 605 additions and 21 deletions
@@ -1,21 +1,26 @@
 import os
+import json
+import uuid
 from typing import Any, Optional, Dict
-from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi import APIRouter, Depends, HTTPException, status, Query
 from pydantic import BaseModel, Field
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import select, func

 from app.dependencies import get_current_user, get_db
 from app.models.user import User
 from app.models.pre_test_template import PreTestTemplate
+from app.models.upload_task import UploadTask
 from app.services.upload_material_service import upload_material_to_platform
-from sqlalchemy import select
+from app.services.upload_queue import upload_queue
 from app.utils.douyinApi import DouyinApi
+from app.utils.id_gen import generate_id


 router = APIRouter(prefix="/upload-material", tags=["上传素材"])


-class UploadTask(BaseModel):
+class UploadTaskRequest(BaseModel):
    advertiser_ids: list[str] = Field(..., description="广告主id数组，支持多条")
    resource_ids: list[str] = Field(..., description="资源id数组（generated_resources表主键）")
    oauth_id: str = Field(..., description="授权表id")
@@ -24,7 +29,7 @@ class UploadTask(BaseModel):


 class BatchUploadRequest(BaseModel):
-    tasks: list[UploadTask] = Field(..., description="批量上传任务列表")
+    tasks: list[UploadTaskRequest] = Field(..., description="批量上传任务列表")


@router.post(
@@ -37,22 +42,6 @@ async def batch_upload_material(
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
 ) -> Any | dict:
-    # try:
-    #     params = {
-    #         "advertiser_id": 1836693172153543,
-    #         "video_ids": ["tos-cn-i-sd07hgqsbj/2db9a53ee3ff4b6f8a2763b3bb463047"],
-    #         "diagnose_config": {"platform": "AD", "external_action": "AD_APP_ACTIVATE"},
-    #     }
-    #     response = await DouyinApi().pre_test_material("0019ecab9b8bc57d964", params)
-    #     return response    
-    # except Exception as e:
-    #     import traceback
-    #     return {
-    #         "code": 500,
-    #         "message": "请求失败",
-    #         "error": str(e),
-    #         "traceback": traceback.format_exc()
-    #     }
    if not req.tasks:
        return {
            "code": 0,
@@ -214,4 +203,85 @@ async def batch_upload_material(
            "total_requested": sum(len(t.resource_ids) * len(t.advertiser_ids) for t in req.tasks),
        },
        "details": all_results,
+    }
+
+
+@router.post(
+    "/async-batch-upload",
+    summary="异步批量上传素材到平台",
+    description="支持批量上传多个授权账户下的资源到素材库，提交后立即返回，后台异步处理",
+)
+async def async_batch_upload_material(
+    req: BatchUploadRequest,
+    current_user: User = Depends(get_current_user),
+    db: AsyncSession = Depends(get_db),
+) -> Any | dict:
+    if not req.tasks:
+        return {
+            "code": 0,
+            "message": "上传任务列表不能为空",
+            "task_ids": [],
+        }
+
+    task_ids = []
+    
+    for task in req.tasks:
+        if not task.advertiser_ids or not task.resource_ids:
+            continue
+        
+        if task.is_pre_test == "1" and not task.pre_test_template:
+            continue
+        
+        for advertiser_id in task.advertiser_ids:
+            for resource_id in task.resource_ids:
+                task_id = generate_id()
+                upload_task = UploadTask(
+                    id = task_id,
+                    user_id=current_user.id,
+                    oauth_id=task.oauth_id,
+                    advertiser_id=advertiser_id,
+                    resource_id=resource_id,
+                    status=1,
+                    note=None,
+                )
+                
+                db.add(upload_task)
+                await upload_queue.enqueue(task_id)
+                task_ids.append(task_id)
+    
+    await db.commit()
+    
+    return {
+        "code": 0,
+        "message": "上传任务已提交",
+        "task_ids": task_ids,
+    }
+
+
+@router.get(
+    "/upload-history",
+    summary="查询上传任务历史",
+    description="查询当前用户的上传任务历史列表",
+)
+async def get_upload_history(
+    page: int = Query(1, ge=1, description="页码"),
+    page_size: int = Query(20, ge=1, le=100, description="每页数量"),
+    status: Optional[int] = Query(None, description="上传状态筛选：1待上传，2上传中，3上传成功，4上传失败"),
+    current_user: User = Depends(get_current_user),
+    db: AsyncSession = Depends(get_db),
+) -> Any | dict:
+    from app.services.upload_material_service import get_upload_history as get_upload_history_service
+    
+    result = await get_upload_history_service(
+        user_id=current_user.id,
+        db=db,
+        page=page,
+        page_size=page_size,
+        status=status,
+    )
+    
+    return {
+        "code": 0,
+        "data": result["data"],
+        "pagination": result["pagination"],
    }
@@ -62,6 +62,11 @@ async def lifespan(app: FastAPI):
    # 启动token刷新定时任务（每5分钟检查一次，小于800秒有效期的token进行刷新）
    from app.tasks.token_refresh_task import token_refresh_scheduler
    token_refresh_task = asyncio.create_task(token_refresh_scheduler())
+
+    # 启动上传任务队列
+    from app.services.upload_queue import upload_queue
+    await upload_queue.recover()
+    upload_queue_task = asyncio.create_task(upload_queue.run())
    
    # 启动时立即同步一次未支付订单
    asyncio.create_task(asyncio.sleep(5))  # 等待5秒后再同步，让系统完全启动
@@ -85,6 +90,8 @@ async def lifespan(app: FastAPI):

    task_queue.stop()
    await queue_task
+    upload_queue.stop()
+    await upload_queue_task
    expiry_task.cancel()
    token_refresh_task.cancel()
    await close_database()
@@ -2,13 +2,14 @@ import os
 import hashlib
 import base64
 import json
-from sqlalchemy import select
+from sqlalchemy import select, func
 from sqlalchemy.ext.asyncio import AsyncSession

 from app.models.generated_resource import GeneratedResource
 from app.models.user_oauth import UserOAuth
 from app.models.resources_material import ResourcesMaterial
 from app.models.pre_test_template import PreTestTemplate
+from app.models.upload_task import UploadTask
 from app.utils.id_gen import generate_id
 from app.utils.douyinApi import DouyinApi

@@ -528,4 +529,92 @@ async def _upload_to_tencent(
        "filename": filename,
        "success": False,
        "error": "腾讯平台上传接口暂未开通",
+    }
+
+
+async def get_upload_history(
+    user_id: str,
+    db: AsyncSession,
+    page: int = 1,
+    page_size: int = 20,
+    status: int | None = None,
+) -> dict:
+    offset = (page - 1) * page_size
+    
+    query = (
+        select(
+            UploadTask,
+            GeneratedResource.resource_type,
+            GeneratedResource.resource_url,
+            GeneratedResource.remote_url,
+            GeneratedResource.storage_type,
+            GeneratedResource.storage_path,
+            GeneratedResource.file_size_bytes,
+            GeneratedResource.model_name,
+        )
+        .outerjoin(
+            GeneratedResource,
+            UploadTask.resource_id == GeneratedResource.id,
+        )
+        .where(
+            UploadTask.user_id == user_id,
+            UploadTask.deleted_at.is_(None),
+        )
+    )
+    
+    if status is not None:
+        query = query.where(UploadTask.status == status)
+    
+    result = await db.execute(
+        query.order_by(UploadTask.created_at.desc())
+        .offset(offset)
+        .limit(page_size)
+    )
+    tasks = result.all()
+    
+    count_query = select(func.count(UploadTask.id)).where(
+        UploadTask.user_id == user_id,
+        UploadTask.deleted_at.is_(None),
+    )
+    
+    if status is not None:
+        count_query = count_query.where(UploadTask.status == status)
+    
+    total_result = await db.execute(count_query)
+    total = total_result.scalar_one()
+    
+    status_map = {
+        1: "待上传",
+        2: "上传中",
+        3: "上传成功",
+        4: "上传失败",
+    }
+    
+    data = []
+    for task, resource_type, resource_url, remote_url, storage_type, storage_path, file_size_bytes, model_name in tasks:
+        data.append({
+            "task_id": task.id,
+            "status": task.status,
+            "status_text": status_map.get(task.status, "未知"),
+            "advertiser_id": task.advertiser_id,
+            "resource_id": task.resource_id,
+            "resource_type": resource_type,
+            "resource_url": resource_url,
+            "remote_url": remote_url,
+            "storage_type": storage_type,
+            "storage_path": storage_path,
+            "file_size_bytes": file_size_bytes,
+            "model_name": model_name,
+            "note": task.note,
+            "created_at": task.created_at,
+            "updated_at": task.updated_at,
+        })
+    
+    return {
+        "data": data,
+        "pagination": {
+            "page": page,
+            "page_size": page_size,
+            "total": total,
+        },
    }
@@ -0,0 +1,418 @@
+import asyncio
+import logging
+from datetime import datetime, timezone
+
+from sqlalchemy import select, update
+from sqlalchemy.orm import selectinload
+
+from app.models.base import async_session
+from app.models.upload_task import UploadTask
+from app.models.generated_resource import GeneratedResource
+from app.models.user_oauth import UserOAuth
+from app.models.resources_material import ResourcesMaterial
+from app.utils.id_gen import generate_id
+from app.utils.douyinApi import DouyinApi
+
+import os
+import hashlib
+
+LOG_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "logs")
+
+os.makedirs(LOG_DIR, exist_ok=True)
+
+logger = logging.getLogger("upload_queue")
+logger.setLevel(logging.INFO)
+
+
+class DailyRotatingFileHandler(logging.FileHandler):
+    def __init__(self, directory, encoding=None):
+        self.directory = directory
+        filename = self._get_log_filename()
+        super().__init__(filename, encoding=encoding)
+    
+    def _get_log_filename(self):
+        return os.path.join(self.directory, f"upload_queue-{datetime.now(timezone.utc).strftime('%Y-%m-%d')}.log")
+    
+    def emit(self, record):
+        current_filename = self._get_log_filename()
+        if self.baseFilename != current_filename:
+            self.close()
+            self.baseFilename = current_filename
+            self.stream = self._open()
+        super().emit(record)
+
+
+if not logger.handlers:
+    handler = DailyRotatingFileHandler(LOG_DIR, encoding="utf-8")
+    handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s", "%Y-%m-%d %H:%M:%S"))
+    logger.addHandler(handler)
+
+douyin_api = DouyinApi()
+
+
+class UploadQueue:
+    def __init__(self):
+        self.queue: asyncio.Queue[str] = asyncio.Queue()
+        self.running = False
+
+    async def enqueue(self, task_id: str):
+        """Add a task to the queue."""
+        await self.queue.put(task_id)
+
+    async def recover(self):
+        """Recover pending tasks from DB on startup."""
+        async with async_session() as db:
+            result = await db.execute(
+                select(UploadTask).where(
+                    UploadTask.status.in_([1, 2]),
+                    UploadTask.deleted_at.is_(None),
+                )
+            )
+            tasks = result.scalars().all()
+            for task in tasks:
+                await self.queue.put(task.id)
+                logger.info(f"Recovered upload task: {task.id}")
+
+    async def run(self):
+        """Main processing loop."""
+        self.running = True
+        logger.info("Upload queue started")
+
+        while self.running:
+            try:
+                task_id = await asyncio.wait_for(self.queue.get(), timeout=5.0)
+            except asyncio.TimeoutError:
+                continue
+
+            try:
+                await self._process(task_id)
+            except Exception as e:
+                logger.error(f"Error processing upload task {task_id}: {e}")
+            finally:
+                self.queue.task_done()
+
+        logger.info("Upload queue stopped")
+
+    async def _process(self, task_id: str):
+        """Process a single upload task."""
+        async with async_session() as db:
+            result = await db.execute(
+                select(UploadTask).where(
+                    UploadTask.id == task_id,
+                    UploadTask.deleted_at.is_(None),
+                ).with_for_update()
+            )
+            task = result.scalar_one_or_none()
+
+            if not task:
+                logger.warning(f"Upload task {task_id} not found or deleted")
+                return
+
+            if task.status == 2:
+                logger.warning(f"Upload task {task_id} is already running")
+                return
+
+            task.status = 2
+            await db.commit()
+
+        try:
+            result = await _upload_single_material(
+                task.id,
+                task.user_id,
+                task.oauth_id,
+                task.advertiser_id,
+                task.resource_id,
+                db=None
+            )
+
+            async with async_session() as db:
+                if result.get("success"):
+                    await db.execute(
+                        update(UploadTask).where(UploadTask.id == task_id).values(
+                            status=3,
+                            note=result.get("message", "上传成功"),
+                        )
+                    )
+                else:
+                    await db.execute(
+                        update(UploadTask).where(UploadTask.id == task_id).values(
+                            status=4,
+                            note=result.get("error", "上传失败"),
+                        )
+                    )
+                await db.commit()
+
+            message = result.get('message') or result.get('error', 'Unknown error')
+            logger.info(f"Upload task {task_id} completed: {'success' if result.get('success') else 'failed'}. Message: {message}")
+
+        except Exception as e:
+            async with async_session() as db:
+                await db.execute(
+                    update(UploadTask).where(UploadTask.id == task_id).values(
+                        status=4,
+                        note=f"上传失败: {str(e)}",
+                    )
+                )
+                await db.commit()
+
+            logger.error(f"Upload task {task_id} failed: {e}")
+
+    def stop(self):
+        """Stop the queue."""
+        self.running = False
+
+
+async def _upload_single_material(
+    task_id: str,
+    user_id: str,
+    oauth_id: str,
+    advertiser_id: str,
+    resource_id: str,
+    db=None
+) -> dict:
+    if db is None:
+        from app.models.base import async_session
+        db = async_session()
+    
+    try:
+        oauth = await db.execute(
+            select(UserOAuth).where(
+                UserOAuth.id == oauth_id,
+                UserOAuth.deleted_at.is_(None),
+                UserOAuth.user_id == user_id,
+            )
+        )
+        oauth = oauth.scalar_one_or_none()
+        if not oauth:
+            return {
+                "success": False,
+                "error": "授权记录不存在"
+            }
+
+        resource = await db.execute(
+            select(GeneratedResource).where(
+                GeneratedResource.id == resource_id,
+                GeneratedResource.user_id == user_id,
+                GeneratedResource.deleted_at.is_(None),
+            )
+        )
+        resource = resource.scalar_one_or_none()
+        
+        if not resource:
+            return {
+                "success": False,
+                "error": f"资源 {resource_id} 不存在或不属于当前用户"
+            }
+
+        resource_type = resource.resource_type
+        storage_path = resource.storage_path
+
+        if resource_type not in ["image", "video"]:
+            return {
+                "success": False,
+                "error": f"不支持的资源类型: {resource_type}，仅支持image和video",
+            }
+
+        if not storage_path:
+            return {
+                "success": False,
+                "error": "资源本地存储路径为空",
+            }
+
+        return await _upload_to_juliang(
+            oauth_id, storage_path, resource_type, advertiser_id, resource, db, user_id
+        )
+    finally:
+        if db is not None:
+            await db.close()
+
+
+async def _upload_to_juliang(
+    oauth_id: str,
+    storage_path: str,
+    resource_type: str,
+    advertiser_id: str,
+    resource: GeneratedResource,
+    db,
+    current_user_id: str,
+) -> dict:
+    filename = os.path.basename(storage_path)
+
+    if resource_type == "image":
+        if resource.file_size_bytes > 5 * 1024 * 1024:
+            return {
+                "resource_id": resource.id,
+                "advertiser_id": advertiser_id,
+                "filename": filename,
+                "success": False,
+                "error": "图片大小超过5M，不支持上传",
+            }
+
+        with open(storage_path, "rb") as f:
+            file_content = f.read()
+            image_signature = hashlib.md5(file_content).hexdigest()
+
+        data = {
+            "advertiser_id": advertiser_id,
+            "upload_type": "UPLOAD_BY_FILE",
+            "image_signature": image_signature,
+            "filename": filename,
+        }
+
+        files = {
+            "image_file": (filename, file_content, "image/png"),
+        }
+
+        response = await douyin_api.upload_image_material(oauth_id, data, files)
+
+        if response["code"] != 0:
+            return {
+                "resource_id": resource.id,
+                "advertiser_id": advertiser_id,
+                "filename": filename,
+                "success": False,
+                "error": response.get("message", "上传失败"),
+            }
+
+        image_id = response["data"]["id"]
+        material_id = response["data"]["material_id"]
+
+        existing = await db.execute(
+            select(ResourcesMaterial).where(
+                ResourcesMaterial.oauth_id == oauth_id,
+                ResourcesMaterial.advertiser_id == advertiser_id,
+                ResourcesMaterial.material_id == str(material_id),
+                ResourcesMaterial.upload_id == str(image_id),
+                ResourcesMaterial.deleted_at.is_(None),
+            )
+        )
+        existing = existing.scalar_one_or_none()
+
+        if existing:
+            return {
+                "resource_id": resource.id,
+                "advertiser_id": advertiser_id,
+                "filename": filename,
+                "success": True,
+                "material_id": str(material_id),
+                "upload_id": str(image_id),
+                "message": "素材已存在，跳过添加",
+            }
+
+        db.add(ResourcesMaterial(
+            id=generate_id(),
+            oauth_id=oauth_id,
+            advertiser_id=advertiser_id,
+            target_table="generated_resources",
+            target_id=resource.id,
+            material_id=str(material_id),
+            upload_id=str(image_id),
+            resource_type=resource_type,
+            user_id=current_user_id,
+        ))
+        await db.commit()
+
+        return {
+            "resource_id": resource.id,
+            "advertiser_id": advertiser_id,
+            "filename": filename,
+            "success": True,
+            "material_id": str(material_id),
+            "upload_id": str(image_id),
+        }
+
+    elif resource_type == "video":
+        if resource.file_size_bytes > 500 * 1024 * 1024:
+            return {
+                "resource_id": resource.id,
+                "advertiser_id": advertiser_id,
+                "filename": filename,
+                "success": False,
+                "error": "视频大小超过500M，不支持上传",
+            }
+
+        with open(storage_path, "rb") as f:
+            file_content = f.read()
+            video_signature = hashlib.md5(file_content).hexdigest()
+
+        data = {
+            "advertiser_id": advertiser_id,
+            "upload_type": "UPLOAD_BY_FILE",
+            "video_signature": video_signature,
+            "filename": filename,
+        }
+
+        files = {
+            "video_file": (filename, file_content, "video/mp4"),
+        }
+
+        response = await douyin_api.upload_video_material(oauth_id, data, files)
+
+        if response["code"] != 0:
+            return {
+                "resource_id": resource.id,
+                "advertiser_id": advertiser_id,
+                "filename": filename,
+                "success": False,
+                "error": response.get("message", "上传失败"),
+            }
+
+        #上传成功，添加视频id，素材id到素材库
+        video_id = response["data"]["video_id"]
+        material_id = response["data"]["material_id"]
+
+        existing = await db.execute(
+            select(ResourcesMaterial).where(
+                ResourcesMaterial.oauth_id == oauth_id,
+                ResourcesMaterial.advertiser_id == advertiser_id,
+                ResourcesMaterial.material_id == str(material_id),
+                ResourcesMaterial.upload_id == str(video_id),
+                ResourcesMaterial.deleted_at.is_(None),
+            )
+        )
+        existing = existing.scalar_one_or_none()
+
+        if existing:
+            return {
+                "resource_id": resource.id,
+                "advertiser_id": advertiser_id,
+                "filename": filename,
+                "success": True,
+                "material_id": str(material_id),
+                "upload_id": str(video_id),
+                "message": "素材已存在，跳过添加",
+            }
+
+        db.add(ResourcesMaterial(
+            id=generate_id(),
+            oauth_id=oauth_id,
+            advertiser_id=advertiser_id,
+            target_table="generated_resources",
+            target_id=resource.id,
+            material_id=str(material_id),
+            upload_id=str(video_id),
+            resource_type=resource_type,
+            user_id=current_user_id,
+        ))
+        await db.commit()
+
+        return {
+            "resource_id": resource.id,
+            "advertiser_id": advertiser_id,
+            "filename": filename,
+            "success": True,
+            "material_id": str(material_id),
+            "upload_id": str(video_id),
+        }
+
+    else:
+        return {
+            "resource_id": resource.id,
+            "advertiser_id": advertiser_id,
+            "filename": filename,
+            "success": False,
+            "error": f"不支持的资源类型: {resource_type}",
+        }
+
+
+upload_queue = UploadQueue()