import asyncio from sqlalchemy import select, update from sqlalchemy.ext.asyncio import AsyncSession from app.models.base import async_session from app.models.resources_material import ResourcesMaterial from app.models.user_oauth import UserOAuth from app.models.pre_test_template import PreTestTemplate from app.utils.logger import get_logger from app.utils.douyinApi import DouyinApi import json logger = get_logger("pre_test_queue", "pre_test_queue") douyin_api = DouyinApi() class PreTestQueue: def __init__(self): self.queue: asyncio.Queue[dict] = asyncio.Queue() self.running = False async def enqueue(self, task_data: dict): """Add a pre-test task to the queue.""" await self.queue.put(task_data) logger.info(f"Enqueued pre-test task: {task_data.get('task_id')}") async def run(self): """Main processing loop.""" self.running = True logger.info("Pre-test queue started") while self.running: try: task_data = await asyncio.wait_for(self.queue.get(), timeout=5.0) except asyncio.TimeoutError: continue try: await self._process(task_data) except Exception as e: logger.error(f"Error processing pre-test task: {e}", exc_info=True) finally: self.queue.task_done() logger.info("Pre-test queue stopped") async def _process(self, task_data: dict): """Process a single pre-test task.""" task_id = task_data.get("task_id") grouped_videos = task_data.get("grouped_videos", {}) pre_test_template_id = task_data.get("pre_test_template_id") logger.info(f"Processing pre-test task: {task_id}") for key, videos in grouped_videos.items(): oauth_id, advertiser_id = key.split("|") video_ids = [v["upload_id"] for v in videos] logger.info(f"Processing {len(video_ids)} videos for oauth_id={oauth_id}, advertiser_id={advertiser_id}") try: async with async_session() as db: await _pre_test_material_batch( oauth_id=oauth_id, advertiser_id=advertiser_id, video_ids=video_ids, pre_test_template_id=pre_test_template_id, db=db, resource_material_ids=[v["id"] for v in videos], ) except Exception as e: logger.error(f"Failed to process pre-test for oauth_id={oauth_id}, advertiser_id={advertiser_id}: {e}", exc_info=True) def stop(self): """Stop the queue.""" self.running = False async def _pre_test_material_batch( oauth_id: str, advertiser_id: str, video_ids: list[str], pre_test_template_id: str, db: AsyncSession, resource_material_ids: list[str], ) -> any: oauth = await db.execute( select(UserOAuth).where( UserOAuth.id == oauth_id, UserOAuth.deleted_at.is_(None), ) ) oauth = oauth.scalar_one_or_none() if not oauth: note = "授权记录不存在" for video_id, resource_material_id in zip(video_ids, resource_material_ids): await _update_material_pre_test_status( db, resource_material_id, task_id=None, status="FAILED", note=note, pre_result=None, pre_test_template_id=pre_test_template_id, ) await db.commit() return {"code": -1, "message": note, "data": {}} pre_test_template = await db.execute( select(PreTestTemplate).where( PreTestTemplate.id == pre_test_template_id, PreTestTemplate.deleted_at.is_(None), PreTestTemplate.user_id == oauth.user_id, ) ) pre_test_template = pre_test_template.scalar_one_or_none() if not pre_test_template: note = f"前测模板 {pre_test_template_id} 不存在" for video_id, resource_material_id in zip(video_ids, resource_material_ids): await _update_material_pre_test_status( db, resource_material_id, task_id=None, status="FAILED", note=note, pre_result=None, pre_test_template_id=pre_test_template_id, ) await db.commit() return {"code": -1, "message": note, "data": {}} diagnose_config = {} if pre_test_template.platform: diagnose_config["platform"] = pre_test_template.platform if pre_test_template.external_action: diagnose_config["external_action"] = pre_test_template.external_action if pre_test_template.cpa_bid: diagnose_config["cpa_bid"] = pre_test_template.cpa_bid if pre_test_template.audience_gender: diagnose_config["audience_gender"] = pre_test_template.audience_gender if pre_test_template.audience_age: diagnose_config["audience_age"] = json.loads(pre_test_template.audience_age) if pre_test_template.audience_region: diagnose_config["audience_region"] = json.loads(pre_test_template.audience_region) if pre_test_template.audience_network: diagnose_config["audience_network"] = json.loads(pre_test_template.audience_network) if pre_test_template.cus_name: diagnose_config["cus_name"] = pre_test_template.cus_name if pre_test_template.pricing_type: diagnose_config["pricing_type"] = pre_test_template.pricing_type if pre_test_template.cost_cap: diagnose_config["cost_cap"] = pre_test_template.cost_cap if pre_test_template.target_cost: diagnose_config["target_cost"] = pre_test_template.target_cost if pre_test_template.nobid: diagnose_config["nobid"] = pre_test_template.nobid if pre_test_template.cpc_bid: diagnose_config["cpc_bid"] = pre_test_template.cpc_bid if pre_test_template.budget: diagnose_config["budget"] = pre_test_template.budget params = { "advertiser_id": int(advertiser_id), "video_ids": video_ids, "diagnose_config": diagnose_config, } response = await douyin_api.pre_test_material(oauth_id, params) code = response.get("code", -1) if code != 0: #记录日志 logger.error(f"前测提交失败, oauth_id={oauth_id}, advertiser_id={advertiser_id}: {json.dumps(response)}") note = response.get("message", "未知错误") for video_id, resource_material_id in zip(video_ids, resource_material_ids): await _update_material_pre_test_status( db, resource_material_id, task_id=None, status="FAILED", note=note, pre_result=None, pre_test_template_id=pre_test_template_id, ) await db.commit() return response data = response.get("data", {}) task_ids = data.get("task_ids", []) fail_video_ids = data.get("fail_video_ids", {}) success_count = 0 for i, (video_id, resource_material_id) in enumerate(zip(video_ids, resource_material_ids)): if video_id in fail_video_ids: fail_info = fail_video_ids[video_id] err_code = fail_info.get("err_code", "") err_message = fail_info.get("err_message", "未知错误") note = f"失败[{err_code}]: {err_message}" await _update_material_pre_test_status( db, resource_material_id, task_id=None, status="FAILED", note=note, pre_result=None, pre_test_template_id=pre_test_template_id, ) else: task_id = str(task_ids[success_count]) if success_count < len(task_ids) else None await _update_material_pre_test_status( db, resource_material_id, task_id=task_id, status="PENDING", note="", pre_result=None, pre_test_template_id=pre_test_template_id, ) success_count += 1 await db.commit() return response async def _update_material_pre_test_status( db: AsyncSession, resource_material_id: str, task_id: str | None, status: str, note: str, pre_result: str | None, pre_test_template_id: str | None, ): await db.execute( update(ResourcesMaterial).where( ResourcesMaterial.id == resource_material_id, ResourcesMaterial.deleted_at.is_(None), ).values( task_id=task_id, status=status, note=note, pre_result=pre_result, pre_test_template_id=pre_test_template_id, ) ) pre_test_queue = PreTestQueue()