knowledge/query_key/set_querys.py

# 读取数据库中的knowledge_content_category表，读出level为1的name和category_id,并将name和query.json中的content_type,step_type,idea_type中的每一项进行组合，输出组合出的结果，并写入到query.json中

import json
import os
import sys
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.mysql_db import MysqlHelper
from loguru import logger

def get_level1_categories():
    """从数据库获取level为1的category名称和ID"""
    try:
        sql = "SELECT id, name FROM knowledge_content_category WHERE level = 1"
        result = MysqlHelper.get_values(sql)
        
        if result:
            # 提取id和name字段的值
            categories = [(row[0], row[1]) for row in result]
            logger.info(f"成功获取到 {len(categories)} 个level为1的category")
            return categories
        else:
            logger.warning("未找到level为1的category数据")
            return []
            
    except Exception as e:
        logger.error(f"获取数据库category数据失败: {e}")
        return []

def load_query_json():
    """加载query.json文件"""
    try:
        query_file_path = os.path.join(os.path.dirname(__file__), "query.json")
        with open(query_file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        logger.info("成功加载query.json文件")
        return data
    except Exception as e:
        logger.error(f"加载query.json文件失败: {e}")
        return None

def generate_combinations(categories, query_data):
    """生成category与query类型的组合"""
    combinations = []
    
    # 获取query.json中的各种类型
    content_formats = [item["name"] for item in query_data.get("content_format", [])]
    stages = [item["name"] for item in query_data.get("stage", [])]
    content_types = [item["name"] for item in query_data.get("content_type", [])]
    
    # 生成所有可能的组合
    for category_id, category_name in categories:
        for content_format in content_formats:
            for stage in stages:
                for content_type in content_types:
                    # 生成拼接的字符串
                    query_word = f"{category_name}{content_format}{stage}{content_type}"
                    
                    # 创建JSON结构
                    combination = {
                        "category_id": category_id,
                        "stage": stage,
                        "content_type": content_type,
                        "content_format": content_format,
                        "query_word": query_word,
                        "status": 0,
                        "create_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    }
                    combinations.append(combination)
    
    logger.info(f"生成了 {len(combinations)} 个组合")
    return combinations

def save_to_database(combinations):
    """将组合结果保存到knowledge_content_query表中"""
    try:
        # 先清空表（可选，根据需求决定）
        # clear_sql = "DELETE FROM knowledge_content_query"
        # MysqlHelper.update_values(clear_sql)
        
        # 插入数据的SQL语句
        insert_sql = """
        INSERT INTO knowledge_content_query 
        (category_id, stage, content_type, content_format, query_word, status, create_time) 
        VALUES (%s, %s, %s, %s, %s, %s, %s)
        """
        
        success_count = 0
        total_count = len(combinations)
        
        logger.info(f"开始插入 {total_count} 条记录到数据库...")
        
        for i, combo in enumerate(combinations, 1):
            params = (
                combo["category_id"],
                combo["stage"],
                combo["content_type"],
                combo["content_format"],
                combo["query_word"],
                combo["status"],
                combo["create_time"]
            )
            
            result = MysqlHelper.update_values(insert_sql, params)
            if result is not None:
                success_count += 1
                # 每插入100条记录输出一次进度
                if i % 100 == 0:
                    logger.info(f"已插入 {i}/{total_count} 条记录")
            else:
                logger.error(f"插入失败: {combo}")
        
        logger.info(f"成功插入 {success_count}/{total_count} 条记录到数据库")
        return success_count
        
    except Exception as e:
        logger.error(f"保存到数据库失败: {e}")
        return 0

def print_combinations(combinations):
    """打印所有组合结果"""
    try:
        print(f"\n=== 生成了 {len(combinations)} 个组合 ===\n")
        
        # 打印前5个组合的JSON结构作为示例
        for i, combo in enumerate(combinations[:5], 1):
            print(f"{i:3d}. {json.dumps(combo, ensure_ascii=False, indent=2)}")
        
        if len(combinations) > 5:
            print(f"... 还有 {len(combinations) - 5} 个组合")
        
        print(f"\n=== 总共 {len(combinations)} 个组合 ===")
        return True
        
    except Exception as e:
        logger.error(f"打印组合结果失败: {e}")
        return False

def main():
    """主函数"""
    logger.info("开始执行set_querys.py")
    
    # 1. 从数据库获取level为1的category
    categories = get_level1_categories()
    if not categories:
        logger.error("无法获取category数据，程序退出")
        return
    
    # 2. 加载query.json文件
    query_data = load_query_json()
    if not query_data:
        logger.error("无法加载query.json文件，程序退出")
        return
    
    # 3. 生成组合
    combinations = generate_combinations(categories, query_data)
    
    # 4. 打印组合结果示例
    if print_combinations(combinations):
        logger.info("组合生成成功")
        
        # 5. 保存到数据库
        logger.info("开始保存到数据库...")
        saved_count = save_to_database(combinations)
        
        if saved_count > 0:
            logger.info(f"程序执行成功完成，共保存 {saved_count} 条记录到数据库")
        else:
            logger.error("数据库保存失败")
    else:
        logger.error("程序执行失败")

if __name__ == "__main__":
    main()