python连接mysql数据库生成测试数据

背景
有时候在做业务逻辑的时候免不了要和数据库进行交互，在大量查询、更新数据的时候可以简单测试一下业务代码的执行效率。

此时就需要保持数据库表中有大量的测试数据。

使用mysql-connector连接数据库可以快速简单的生成大量的测试数据，方便后续自己测试使用。

import mysql.connector
from mysql.connector import Error
import random
import string
from datetime import timedelta, datetime

# ！加载依赖 pip install mysql-connector


def generate_random_string(length):
    """生成随机字符串"""
    letters = string.ascii_lowercase
    return "".join(random.choice(letters) for i in range(length))


def generate_random_date(day=3):
    """生成day天前的时间戳"""
    timedelta_days_ago = datetime.now() - timedelta(days=day)

    # 格式化为字符串
    formatted_time = timedelta_days_ago.strftime("%Y-%m-%d %H:%M:%S")

    return formatted_time


# ret = generate_random_string(12)
# print(ret) # ydysiivvrwgw

# ret = generate_random_date(2)
# print(ret) # 2024-07-09 17:16:55


def insert_data(connection, num_records):
    """插入指定num_records数量的记录"""
    try:
        cursor = connection.cursor()
        sql = "INSERT INTO your_table (id, company_id, created_time) VALUES (%s, %s, %s)"
        # records = [(generate_random_string(10), random.randint(18, 65)) for _ in range(num_records)]
        records = [(i, random.randint(1, 6), generate_random_date()) for i in range(num_records)]
        cursor.executemany(sql, records)
        connection.commit()
        print(f"{cursor.rowcount} records inserted successfully into table")
    except Error as e:
        print(f"Failed to insert records: {e}")
    finally:
        if connection.is_connected():
            cursor.close()


def main():
    try:
        connection = mysql.connector.connect(
            host="127.0.0.1", database="btsaas", user="root", password="123456", port=3306
        )
        if connection.is_connected():
            db_info = connection.get_server_info()
            print("Connected to MySQL Server version ", db_info)
            insert_data(connection, 1000000)  # 插入数据
    except Error as e:
        print("Error while connecting to MySQL", e)
    finally:
        if connection.is_connected():
            connection.close()
            print("MySQL connection is closed")


if __name__ == "__main__":
    main()