video-gen/video-gen-api/app/tasks/pre_test_result_task.py

from datetime import datetime, timezone
import asyncio
import os
import logging
import json

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.utils.douyinApi import DouyinApi

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"pre_test_result_task-{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()

#获取前测结果并更新数据库，计划任务，每2分钟执行一次
async def poll_pre_test_results():
    """每2分钟轮询前测结果并更新数据库"""
    logger.info("Pre-test result polling task started")
    
    while True:
        try:
            await process_pending_pre_tests()
            await asyncio.sleep(120)
        
        except asyncio.CancelledError:
            logger.info("Pre-test result polling task cancelled")
            break
        except Exception as e:
            logger.error(f"Error in poll_pre_test_results: {e}")
            await asyncio.sleep(120)


async def process_pending_pre_tests():
    """处理所有待查询的前测任务"""
    async with async_session() as db:
        result = await db.execute(
            select(ResourcesMaterial).where(
                ResourcesMaterial.status == "PENDING",
                ResourcesMaterial.task_id.is_not(None),
                ResourcesMaterial.deleted_at.is_(None),
            )
        )
        pending_materials = result.scalars().all()
        
        if not pending_materials:
            return
        
        logger.info(f"Found {len(pending_materials)} pending pre-test tasks to process")
        
        for material in pending_materials:
            try:
                await update_single_pre_test_result(db, material)
            except Exception as e:
                logger.error(f"Error processing pre-test for material {material.id}: {e}")


async def update_single_pre_test_result(db: AsyncSession, material: ResourcesMaterial):
    """更新单个素材的前测结果"""
    oauth = await db.execute(
        select(UserOAuth).where(
            UserOAuth.id == material.oauth_id,
            UserOAuth.deleted_at.is_(None),
        )
    )
    oauth = oauth.scalar_one_or_none()
    
    if not oauth:
        logger.error(f"OAuth record not found for material {material.id}")
        await db.execute(
            update(ResourcesMaterial).where(ResourcesMaterial.id == material.id).values(
                status="FAILED",
                note="授权记录不存在",
            )
        )
        await db.commit()
        return
    
    params = {
        "advertiser_id": int(material.advertiser_id),
        "task_ids": json.dumps([int(material.task_id)]),
    }
    
    try:
        response = await douyin_api.get_material_pre_test_result(oauth.id, params)
    except Exception as e:
        logger.error(f"Failed to get pre-test result for task {material.task_id}: {e}")
        return
    
    code = response.get("code", -1)
    if code != 0:
        logger.error(f"API error for task {material.task_id}: {response.get('message', 'Unknown error')}")
        return
    
    data = response.get("data", {})
    task_details = data.get("task_list", [])
    
    if not task_details:
        return
    
    task_detail = task_details[0]
    status = task_detail.get("status") 
    pre_result = {
        "video_id": task_detail.get("video_id") or None,
        "advertiser_id": task_detail.get("advertiser_id") or None,
        "material_id": task_detail.get("material_id") or None,
        "is_ad_high_quality_material": task_detail.get("is_ad_high_quality_material") or None,
        "is_ecp_high_quality_material": task_detail.get("is_ecp_high_quality_material") or None,
        "is_inefficient_material": task_detail.get("is_inefficient_material") or None,
        "is_first_publish_material": task_detail.get("is_first_publish_material") or None,
        "not_ad_high_quality_reason": task_detail.get("not_ad_high_quality_reason") or None,
        "not_ecp_high_quality_reason": task_detail.get("not_ecp_high_quality_reason") or None,
        "is_local_high_quality_material": task_detail.get("is_local_high_quality_material") or None,
    }
    
    await db.execute(
        update(ResourcesMaterial).where(ResourcesMaterial.id == material.id).values(
            status=status,
            pre_result=json.dumps(pre_result, ensure_ascii=False),
        )
    )
    await db.commit()