386 lines
12 KiB
Python
386 lines
12 KiB
Python
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
|
|
|
|
logger = logging.getLogger("upload_queue")
|
|
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()
|
|
|
|
logger.info(f"Upload task {task_id} completed: {'success' if result.get('success') else 'failed'}")
|
|
|
|
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", "上传失败"),
|
|
}
|
|
|
|
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() |