LongArticlesJob/tasks/update_article_info_from_aigc.py

"""
@author: luojunhui
"""

import json
from typing import List, Dict

from pymysql.cursors import DictCursor
from tqdm import tqdm

from applications.const import ArticleCollectorConst
from applications.db import DatabaseConnector
from applications.functions import Functions
from config import denet_config, long_articles_config

empty_dict = {}
const = ArticleCollectorConst()
functions = Functions()


class UpdateArticleInfoFromAIGC(object):
    """
    从aigc获取文章信息
    """

    def __init__(self):
        self.aigc_db_client = DatabaseConnector(db_config=denet_config)
        self.long_articles_db_client = DatabaseConnector(db_config=long_articles_config)
        self.aigc_db_client.connect()
        self.long_articles_db_client.connect()

    def get_published_articles(self) -> List[Dict]:
        """
        获取当天发布文章的List
        """
        sql = f"""
            SELECT trace_id, push_type
            FROM long_articles_published_trace_id
            WHERE create_timestamp> UNIX_TIMESTAMP(DATE_SUB(CURDATE(), INTERVAL 1 DAY)) AND status = {const.INIT_STATUS};
        """
        article_list = self.long_articles_db_client.fetch(sql, cursor_type=DictCursor)
        return article_list

    def get_article_info_from_aigc(self, trace_id: str) -> Dict:
        """
        从aigc获取发布结果
        """
        sql = f"""
            SELECT t2.crawler_channel_content_id, t2.publish_stage_url, t2.publish_timestamp, t1.result_data
            from publish_content_miniprogram t1
            join publish_content t2 on t1.publish_content_id = t2.id
            where t1.trace_id = '{trace_id}' and t2.status = {const.PUBLISHED_STATUS};
        """
        article_info = self.aigc_db_client.fetch(sql, cursor_type=DictCursor)
        if article_info:
            return article_info[0]
        else:
            return empty_dict

    def update_each_article(self, article: Dict):
        """
        更新每个文章的信息
        """
        trace_id = article["trace_id"]
        push_type = article["push_type"]

        article_info = self.get_article_info_from_aigc(trace_id)
        if article_info:
            channel_content_id = article_info["crawler_channel_content_id"]
            published_url = article_info["publish_stage_url"]
            publish_timestamp = int(article_info["publish_timestamp"] / 1000)
            result_data = json.loads(article_info["result_data"])
            root_source_id_list = [
                functions.extract_path(item["productionPath"])["root_source_id"] for item in result_data
            ]
            if published_url:
                status = const.SUCCESS_STATUS
            else:
                if push_type == const.BULK_AUTO_PUSH:
                    status = const.INIT_STATUS
                else:
                    status = const.SUCCESS_STATUS
            update_sql = f"""
                UPDATE long_articles_published_trace_id
                SET published_url = %s, status = %s, publish_timestamp = %s, crawler_channel_content_id = %s, root_source_id_list = %s
                WHERE trace_id = %s;
            """
            self.long_articles_db_client.save(
                query=update_sql,
                params=(published_url, status, publish_timestamp, channel_content_id, json.dumps(root_source_id_list), trace_id)
            )
        else:
            update_sql = f"""
                UPDATE long_articles_published_trace_id
                SET status = %s 
                WHERE trace_id = %s;
            """
            self.long_articles_db_client.save(
                query=update_sql, params=(const.FAIL_STATUS, trace_id)
            )

    def deal(self):
        """
        main function
        """
        article_list = self.get_published_articles()
        for article in tqdm(article_list, desc="更新文章信息"):
            try:
                self.update_each_article(article)
            except Exception as e:
                print(e)


if __name__ == "__main__":
    u = UpdateArticleInfoFromAIGC()
    u.deal()