根据示例数据生成更多测试数据脚本

111 阅读1分钟
import time
import mysql.connector
import random


# 建立 MySQL 连接
def create_mysql_connection():
    conn = mysql.connector.connect(
        host="xxxx",
        port="3306",
        user="xxx",
        password="xxx",
        database="xxxx"
    )
    return conn


# 读取表中已有数据
def read_existing_data(conn, table_name):
    cursor = conn.cursor()
    query = f"SELECT * FROM {table_name} LIMIT 100"
    cursor.execute(query)
    existing_data = cursor.fetchall()
    cursor.close()

    # 获取字段名称
    column_names = [column[0] for column in cursor.description]

    # 构建带有字段名称的数据列表
    data_with_column_names = [dict(zip(column_names, data)) for data in existing_data]

    return data_with_column_names


# 推测生成新数据
def generate_new_data(existing_data):
    new_data = {}  # 存储新生成的数据,使用字典类型

    # 获取字段名称列表
    column_names = existing_data[0]

    # 遍历字段名称列表
    for column_name in column_names:
        # 获取已有数据中该字段的值列表
        existing_values = [data[column_name] for data in existing_data[1:]]

        # 检查字段名是否为 "id"
        if column_name == "id":
            # 使用当前时间戳作为新值
            new_value = int(time.time())
        else:
            # 从已有值列表中随机选择一个值作为新的字段值
            new_value = random.choice(existing_values)

        # 将字段名和字段值添加到新数据字典中
        new_data[column_name] = new_value

    return new_data


# 插入新数据到表中
def insert_new_data(conn, table_name, new_data):
    cursor = conn.cursor()

    # 获取字段列表(排除主键字段)
    cursor.execute(f"SHOW COLUMNS FROM {table_name}")
    columns = [column[0] for column in cursor.fetchall() if column[3] != 'PRI']
    cursor.close()

    # 检查字段数量与数据项数量是否一致
    if len(columns) != len(new_data)-1:
        print("字段数量与数据项数量不一致", len(columns), " ", len(new_data))
        return

    # 构建插入语句
    column_names = ', '.join(columns)
    value_placeholders = ', '.join(['%s'] * len(columns))
    query = f"INSERT INTO {table_name} ({column_names}) VALUES ({value_placeholders})"

    # 将字段值按照字段名的顺序构建成列表
    values = [new_data[column] for column in columns]

    # 执行插入操作
    cursor = conn.cursor()
    cursor.execute(query, values)
    conn.commit()
    cursor.close()


# 主函数
def main(table_name):
    # 建立 MySQL 连接
    conn = create_mysql_connection()

    # 读取已有数据
    existing_data = read_existing_data(conn, table_name)

    # 推测生成新数据
    new_data = generate_new_data(existing_data)

    # 插入新数据到表中
    insert_new_data(conn, table_name, new_data)

    # 关闭连接
    conn.close()

if __name__ == "__main__":
    table_name = "xxxx"
    main(table_name)