aitsc/direct_insert_templates.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
直接插入模板数据到腾讯云数据库
不依赖Flask应用，直接从文件读取模板数据
"""
import pymysql
import re
import json

def extract_templates_from_file():
    """直接从文件提取模板数据"""
    print("📖 正在读取模板文件...")
    
    try:
        with open('src/flask_prompt_master/promptsTemplates.py', 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 查找templates列表的开始和结束
        start_pattern = r'templates\s*=\s*\['
        end_pattern = r'\]\s*$'
        
        start_match = re.search(start_pattern, content)
        if not start_match:
            print("❌ 无法找到templates列表")
            return []
        
        # 找到templates列表的开始位置
        start_pos = start_match.end()
        
        # 从开始位置查找对应的结束括号
        bracket_count = 1
        pos = start_pos
        while pos < len(content) and bracket_count > 0:
            if content[pos] == '[':
                bracket_count += 1
            elif content[pos] == ']':
                bracket_count -= 1
            pos += 1
        
        if bracket_count != 0:
            print("❌ 无法找到templates列表的结束位置")
            return []
        
        # 提取templates列表的字符串
        templates_str = content[start_pos-1:pos]
        
        # 使用eval来解析Python列表（注意：这在生产环境中不推荐，但这里是为了简化）
        try:
            templates = eval(templates_str)
            print(f"✅ 成功提取 {len(templates)} 个模板")
            return templates
        except Exception as e:
            print(f"❌ 解析模板数据失败: {str(e)}")
            return []
            
    except Exception as e:
        print(f"❌ 读取文件失败: {str(e)}")
        return []

def insert_templates_to_tencent():
    """插入模板数据到腾讯云数据库"""
    print("🚀 开始插入模板数据到腾讯云数据库...")
    
    # 腾讯云数据库配置
    config = {
        'host': 'gz-cynosdbmysql-grp-d26pzce5.sql.tencentcdb.com',
        'port': 24936,
        'user': 'root',
        'password': '!Rjb12191',
        'database': 'pro_db',
        'charset': 'utf8mb4'
    }
    
    try:
        # 连接数据库
        print("🔗 连接到腾讯云数据库...")
        conn = pymysql.connect(**config)
        cursor = conn.cursor()
        print("✅ 数据库连接成功")
        
        # 获取模板数据
        templates = extract_templates_from_file()
        if not templates:
            print("❌ 无法获取模板数据，退出")
            return
        
        # 清空现有数据
        print("🗑️ 清空现有模板数据...")
        cursor.execute("TRUNCATE TABLE prompt_template")
        print("✅ 现有数据已清空")
        
        # 插入所有模板数据
        print("📝 开始插入所有模板数据...")
        sql = """
        INSERT INTO prompt_template 
        (name, description, category, industry, profession, sub_category, system_prompt, is_default) 
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
        """
        
        success_count = 0
        error_count = 0
        
        for i, template in enumerate(templates, 1):
            try:
                cursor.execute(sql, (
                    template.get('name', ''),
                    template.get('description', ''),
                    template.get('category', ''),
                    template.get('industry', ''),
                    template.get('profession', ''),
                    template.get('sub_category', ''),
                    template.get('system_prompt', ''),
                    template.get('is_default', False)
                ))
                success_count += 1
                
                # 每插入20个模板显示一次进度
                if i % 20 == 0:
                    print(f"📈 已插入 {i}/{len(templates)} 个模板...")
                    
            except Exception as e:
                print(f"⚠️ 插入模板 '{template.get('name', 'Unknown')}' 失败: {str(e)}")
                error_count += 1
        
        # 提交事务
        conn.commit()
        
        print("\n" + "="*50)
        print("🎉 模板数据插入完成！")
        print(f"✅ 成功插入: {success_count} 个模板")
        if error_count > 0:
            print(f"⚠️ 插入失败: {error_count} 个模板")
        print(f"📊 总计模板: {len(templates)} 个")
        print("="*50)
        
        # 验证插入结果
        cursor.execute("SELECT COUNT(*) FROM prompt_template")
        final_count = cursor.fetchone()[0]
        print(f"🔍 数据库中的模板总数: {final_count}")
        
    except Exception as e:
        print(f"❌ 插入模板数据失败: {str(e)}")
        import traceback
        traceback.print_exc()
        if 'conn' in locals():
            conn.rollback()
    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'conn' in locals():
            conn.close()

def main():
    """主函数"""
    print("=" * 60)
    print("🔄 直接插入模板数据工具")
    print("=" * 60)
    print("⚠️ 警告：此操作将清空现有的所有模板数据！")
    
    # 确认操作
    confirm = input("\n是否继续？(y/N): ").strip().lower()
    if confirm not in ['y', 'yes', '是']:
        print("❌ 操作已取消")
        return
    
    insert_templates_to_tencent()

if __name__ == '__main__':
    main()