RAGflow/python_sql/add_sql_final.py

import mysql.connector
import uuid
import base64
import json
from datetime import datetime
from Cryptodome.PublicKey import RSA
from Cryptodome.Cipher import PKCS1_v1_5
from werkzeug.security import generate_password_hash

# 数据库连接配置
db_config = {
    "host": "localhost",
    "port": 5455,
    "user": "root",
    "password": "infini_rag_flow",
    "database": "rag_flow",
}

# 生成随机的 UUID 作为 id
def generate_uuid():
    return str(uuid.uuid4()).replace("-", "")

# RSA 加密密码
def rsa_psw(password: str) -> str:
    pub_key = """-----BEGIN PUBLIC KEY-----
    MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArq9XTUSeYr2+N1h3Afl/z8Dse/2yD0ZGrKwx+EEEcdsBLca9Ynmx3nIB5obmLlSfmskLpBo0UACBmB5rEjBp2Q2f3AG3Hjd4B+gNCG6BDaawuDlgANIhGnaTLrIqWrrcm4EMzJOnAOI1fgzJRsOOUEfaS318Eq9OVO3apEyCCt0lOQK6PuksduOjVxtltDav+guVAA068NrPYmRNabVKRNLJpL8w4D44sfth5RvZ3q9t+6RTArpEtc5sh5ChzvqPOzKGMXW83C95TxmXqpbK6olN4RevSfVjEAgCydH6HN6OhtOQEcnrU97r9H0iZOWwbw3pVrZiUkuRD1R56Wzs2wIDAQAB
    -----END PUBLIC KEY-----"""

    rsa_key = RSA.import_key(pub_key)
    cipher = PKCS1_v1_5.new(rsa_key)
    encrypted_data = cipher.encrypt(base64.b64encode(password.encode()))
    return base64.b64encode(encrypted_data).decode()

# 加密密码
def encrypt_password(raw_password: str) -> str:
    base64_password = base64.b64encode(raw_password.encode()).decode()
    encrypted_password = rsa_psw(base64_password)
    return generate_password_hash(base64_password)

# 处理批量注册
def batch_register_students():
    # 从 JSON 文件加载学生数据
    with open("add.json", "r", encoding="utf-8") as json_file:
        student_groups = json.load(json_file)

    try:
        # 建立数据库连接
        conn = mysql.connector.connect(**db_config)
        cursor = conn.cursor()

        for group in student_groups:
            tenant_id = group["tenant_id"]
            student_ids = group["student_id"]

            for student_id in student_ids:
                # 生成 ID 和时间戳
                user_id = generate_uuid()
                create_time = 1741361741738
                create_date = datetime.strptime("2025-03-07 23:35:41", "%Y-%m-%d %H:%M:%S")
                update_time = 1741416354403
                update_date = datetime.strptime("2025-03-08 14:45:54", "%Y-%m-%d %H:%M:%S")

                # 用户信息
                student_email = student_id + "@xidian.cn"
                raw_password = student_id
                hash_encrypted = encrypt_password(raw_password)

                # 插入 user 表数据
                user_insert_query = """
                INSERT INTO user (
                    id, create_time, create_date, update_time, update_date, access_token,
                    nickname, password, email, avatar, language, color_schema, timezone,
                    last_login_time, is_authenticated, is_active, is_anonymous, login_channel,
                    status, is_superuser
                ) VALUES (
                    %s, %s, %s, %s, %s, %s,
                    %s, %s, %s, %s, %s, %s, %s,
                    %s, %s, %s, %s, %s,
                    %s, %s
                )
                """
                user_data = (
                    user_id, create_time, create_date, update_time, update_date, None,
                    student_id, hash_encrypted, student_email, None, "Chinese", "Bright", "UTC+8 Asia/Shanghai",
                    create_date, 1, 1, 0, "password",
                    1, 0
                )
                cursor.execute(user_insert_query, user_data)

                # 插入 tenant 表数据
                tenant_insert_query = """
                INSERT INTO tenant (
                    id, create_time, create_date, update_time, update_date, name,
                    public_key, llm_id, embd_id, asr_id, img2txt_id, rerank_id, tts_id,
                    parser_ids, credit, status
                ) VALUES (
                    %s, %s, %s, %s, %s, %s,
                    %s, %s, %s, %s, %s, %s, %s,
                    %s, %s, %s
                )
                """
                tenant_data = (
                    user_id, create_time, create_date, update_time, update_date, student_id + "'s Kingdom",
                    None, "deepseek-r1:1.5b@Ollama", "BAAI/bge-large-zh-v1.5@BAAI", "", "", "", None,
                    "naive:General,qa:Q&A,resume:Resume,manual:Manual,table:Table,paper:Paper,book:Book,laws:Laws,presentation:Presentation,picture:Picture,one:One,audio:Audio,email:Email,tag:Tag",
                    512, 1
                )
                cursor.execute(tenant_insert_query, tenant_data)

                # 插入 user_tenant 表的第一条记录
                user_tenant_insert_query = """
                INSERT INTO user_tenant (
                    id, create_time, create_date, update_time, update_date, user_id,
                    tenant_id, role, invited_by, status
                ) VALUES (
                    %s, %s, %s, %s, %s, %s,
                    %s, %s, %s, %s
                )
                """
                user_tenant_data_owner = (
                    generate_uuid(), create_time, create_date, update_time, update_date, user_id,
                    user_id, "owner", user_id, 1
                )
                cursor.execute(user_tenant_insert_query, user_tenant_data_owner)

                # 插入 user_tenant 表的第二条记录
                user_tenant_data_normal = (
                    generate_uuid(), create_time, create_date, update_time, update_date, user_id,
                    tenant_id, "normal", tenant_id, 1
                )
                cursor.execute(user_tenant_insert_query, user_tenant_data_normal)

                # 插入 tenant_llm 表数据
                tenant_llm_insert_query = """
                INSERT INTO tenant_llm (
                    create_time, create_date, update_time, update_date, tenant_id,
                    llm_factory, model_type, llm_name, api_key, api_base, max_tokens, used_tokens
                ) VALUES (
                    %s, %s, %s, %s, %s,
                    %s, %s, %s, %s, %s, %s, %s
                )
                """
                tenant_llm_data = (
                    create_time, create_date, update_time, update_date, user_id,
                    "Ollama", "chat", "deepseek-r1:1.5b", "xxxxxxxxxxxxxxx", "http://10.195.140.47:11434", 88888888, 0
                )
                cursor.execute(tenant_llm_insert_query, tenant_llm_data)

        # 提交事务
        conn.commit()

        print("批量用户数据添加成功！")

    except mysql.connector.Error as err:
        print(f"数据库错误: {err}")

    finally:
        if conn.is_connected():
            cursor.close()
            conn.close()
            print("数据库连接已关闭。")


if __name__ == '__main__':
    batch_register_students()