from __future__ import print_function
import string
import argparse
import random
import threading
import time
import sys
from contextlib import contextmanager
import pymysql
DB_NAME = 'test_insert_data_db'
TABLE_NAME = 'test_insert_data_table'
CREATE_TABLE_STATEMENT = """create table {0} (id int(10) NOT NULL AUTO_INCREMENT,
name varchar(255) NOT NULL, datetime varchar(255) NOT NULL,
PRIMARY KEY (`id`))""".format(TABLE_NAME)
def _argparse():
parser = argparse.ArgumentParser(description='benchmark tool for MySQL database')
parser.add_argument('--host', action='store', dest='host',
required=True, help='connect to host')
parser.add_argument('--user', action='store', dest='user',
required=True, help='user for login')
parser.add_argument('--password', action='store', dest='password',
required=True, help='password to use when connecting to server')
parser.add_argument('--port', action='store', dest='port', default=3306,
type=int, help='port number to use for connection or 3306 for default')
parser.add_argument('--thread_size', action='store', dest='thread_size',
default=5, type=int, help='how much connection for database usage')
parser.add_argument('--row_size', action='store', dest='row_size',
default=5000, type=int, help='how much rows')
parser.add_argument('-v', '--version', action='version', version='%(prog)s 0.1')
return parser.parse_args()
@contextmanager
def get_conn(**kwargs):
conn = pymysql.connect(**kwargs)
try:
yield conn
finally:
conn.close()
def create_db_and_table(conn):
with conn as cur:
for sql in ["drop database if exists {0}".format(DB_NAME),
"create database {0}".format(DB_NAME),
"use {0}".format(DB_NAME),
CREATE_TABLE_STATEMENT]:
print(sql)
cur.execute(sql)
def random_string(length=10):
s = string.ascii_letters + string.digits
return "".join(random.sample(s, length))
def add_row(cursor):
SQL_FORMAT = "INSERT INTO `{0}`(name, datetime) values('{1}', {2})"
sql = SQL_FORMAT.format(TABLE_NAME, random_string(), '\'' + time.strftime('%Y-%m-%d %H:%M:%S') + '\'')
cursor.execute(sql)
def insert_data(conn_args, row_size):
with get_conn(**conn_args) as conn:
with conn as c:
c.execute('use {0}'.format(DB_NAME))
with conn as c:
for i in range(row_size):
add_row(c)
conn.commit()
def main():
parser = _argparse()
conn_args = dict(host=parser.host, user=parser.user,
password=parser.password, port=parser.port)
with get_conn(**conn_args) as conn:
create_db_and_table(conn)
threads = []
for i in range(parser.thread_size):
t = threading.Thread(target=insert_data, args=(conn_args, parser.row_size))
threads.append(t)
t.start()
for t in threads:
t.join()
if __name__ == '__main__':
sys.argv.extend(['--host', '127.0.0.1', '--user', 'root', '--password', '123456'])
main()
打造一个mysql压测工具
最新推荐文章于 2024-05-22 17:27:01 发布