aitsc/test_db.py

from flask_prompt_master import create_app, db
import pymysql
from flask_prompt_master.init_db import templates as init_templates
from flask_prompt_master.templates.prompts import templates as prompt_templates

def insert_all_templates():
    """向 prompt_template 表插入所有模板数据"""
    try:
        # 合并两个模板列表
        all_templates = init_templates + prompt_templates
        print(f"准备插入的模板总数: {len(all_templates)}")
        
        # 连接MySQL数据库
        conn = pymysql.connect(
            host='localhost',
            user='root',
            password='123456',
            database='food_db',
            charset='utf8mb4'
        )
        cursor = conn.cursor()
        
        # 先清空表
        cursor.execute("TRUNCATE TABLE prompt_template")
        
        # SQL 插入语句
        check_sql = "SELECT COUNT(*) FROM prompt_template WHERE name = %s"
        insert_sql = """
        INSERT INTO prompt_template 
        (name, description, category, industry, profession, sub_category, system_prompt) 
        VALUES (%(name)s, %(description)s, %(category)s, %(industry)s, %(profession)s, 
                %(sub_category)s, %(system_prompt)s)
        """
        
        # 遍历所有模板数据并插入
        success_count = 0
        duplicate_count = 0
        error_count = 0
        
        for template in all_templates:
            try:
                # 检查模板名称是否已存在
                cursor.execute(check_sql, (template['name'],))
                exists = cursor.fetchone()[0] > 0
                
                if exists:
                    print(f"模板已存在，跳过: {template['name']}")
                    duplicate_count += 1
                    continue
                
                # 准备模板数据
                template_data = {
                    'name': template['name'],
                    'description': template['description'],
                    'category': template.get('category', ''),
                    'industry': template.get('industry', ''),
                    'profession': template.get('profession', ''),
                    'sub_category': template.get('sub_category', ''),
                    'system_prompt': template['system_prompt']
                }
                
                # 执行插入
                cursor.execute(insert_sql, template_data)
                success_count += 1
                
                print(f"成功插入模板: {template['name']}")
                
            except Exception as e:
                print(f"插入模板 {template['name']} 失败:")
                print(f"错误类型: {type(e).__name__}")
                print(f"错误信息: {str(e)}")
                error_count += 1
                continue
        
        # 提交事务
        conn.commit()
        
        print("\n=== 数据插入完成 ===")
        print(f"成功插入: {success_count} 个模板")
        print(f"重复跳过: {duplicate_count} 个模板")
        print(f"插入失败: {error_count} 个模板")
        print(f"总计模板: {len(all_templates)} 个")
        print("===================")
            
    except Exception as e:
        print(f"数据库连接失败: {str(e)}")
        if 'conn' in locals():
            conn.rollback()
    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'conn' in locals():
            conn.close()

if __name__ == '__main__':
    insert_all_templates()