Files

246 lines
8.8 KiB
Python

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()