aigc-test/result_save.py

import mysql.connector
import json

# 配置数据库连接参数
db_config = {
    'host': 'rm-bp19uc56sud25ag4o.mysql.rds.aliyuncs.com',
    'database': 'longvideo',
    'port': 3306,
    'user': 'wx2016_longvideo',
    'password': 'wx2016_longvideoP@assword1234',
}

json_field_names = ['key_words', 'search_keys', 'extra_keys', 'category']

normal_field_names = ['tone', 'target_audience',
                      'target_age', 'target_gender', 'address', 'theme']


def insert_content(gpt_res):
    """ 连接MySQL数据库并插入一行数据 """
    try:
        # 连接MySQL数据库
        conn = mysql.connector.connect(
            host=db_config['host'],
            database=db_config['database'],
            user=db_config['user'],
            password=db_config['password'],
        )

        if conn.is_connected():
            print('成功连接到数据库')
            cursor = conn.cursor()

            insert_batch = []

            # 插入数据的SQL语句
            sql = """
            INSERT INTO video_content_mapping (video_id, tag, tag_type)
            VALUES (%s, %s, %s)
            """

            video_id = gpt_res.get('video_id', '')
            for field_name in json_field_names:
                # 获取对应的JSON字符串
                tags = gpt_res.get(field_name, '')
                # 判断是否是json字符串
                if not tags or not isinstance(tags, list):
                    continue
                # 构建批量插入的参数
                for tag in tags:
                    insert_batch.append((video_id, tag, field_name))

            for field_name in normal_field_names:
                # 获取对应的字段值
                value = gpt_res.get(field_name, '')
                # 构建批量插入的参数
                insert_batch.append((video_id, value, field_name))

            # 执行批量插入操作
            cursor.executemany(sql, insert_batch)
            print(f"Inserting records {len(insert_batch)} rows...")
            insert_batch.clear()

            # 提交事务
            conn.commit()

            # 关闭游标和连接
            cursor.close()
            conn.close()
            print('数据库连接已关闭')
    except mysql.connector.Error as e:
        print('数据库连接或操作出错:', e)