[python运维] 使用python3制作一个mysql压测小工具！

0x01 argparse模块

argparse是Python 标准库中推荐的命令行解析模块，他可以实现给出参数提示和接收用户输入的参数，类似linux的命令行功能：

[root@yunweisn ~]# grep
Usage: grep [OPTION]... PATTERN [FILE]...
Try 'grep --help' for more information.
[root@yunweisn ~]# grep --help
Usage: grep [OPTION]... PATTERN [FILE]...
Search for PATTERN in each FILE or standard input.
PATTERN is, by default, a basic regular expression (BRE).
Example: grep -i 'hello world' menu.h main.c

Regexp selection and interpretation:
  -E, --extended-regexp     PATTERN is an extended regular expression (ERE)
  -F, --fixed-strings       PATTERN is a set of newline-separated fixed strings

argparse快速入门：

#!/usr/bin/python3
import argparse
parse = argparse.ArgumentParser(description='this is test script') # 实例化
# description 描述信息
parse.add_argument('-t','--test',dest='teststr',action='store',
                    type=str,default='testing...',help='test help string') # 增加一个参数
# -t 短参数
# --test 长参数
# dest 后续引用的变量名，不写默认是长参数名称
# action有6种，不指定默认为store
# type 默认为str，常见的还有int
# default 参数的默认值
# help 用户使用help命令时显示的信息
args = parse.parse_args() # 接收用户输入的参数
print(args.teststr)


# 运行1
[root@yunweisn testdemo]# python3 testdemo01.py --help
usage: testdemo01.py [-h] [-t TESTSTR]

this is test script # 描述信息

optional arguments:
  -h, --help            show this help message and exit # 默认有个help参数
  -t TESTSTR, --test TESTSTR # 参数名称
                        test help string # 参数帮助


# 运行2
[root@yunweisn testdemo]# python3 testdemo01.py
testing...
[root@yunweisn testdemo]# python3 testdemo01.py -t 123
123

argparse有6种action

1. store 保存参数值，可能会先将参数值转换成另一个数据类型。若没有显式指定动作，则默认为该动作。

2. store_const 保存一个被定义为参数规格一部分的值，而不是一个来自参数解析而来的值。这通常用于实现非布尔值的命令行标记。

3. store_ture/store_false 保存相应的布尔值。这两个动作被用于实现布尔开关。

4. append 将值保存到一个列表中。若参数重复出现，则保存多个值。

5. append_const 将一个定义在参数规格中的值保存到一个列表中。

6. version 打印关于程序的版本信息，然后退出

---

0x02 来打造一个类mysql命令行客户端工具

先看看mysql命令默认是如何提示的？

[root@yunweisn ~]# mysql
ERROR 1045 (28000): Access denied for user 'root'@'localhost' (using password: NO)

# 这里可见执行mysql的时候，有两个默认参数 
# -u 用户名
# -h 主机名

python实现类似功能如下：

import argparse
import commands

def main():
    parse = argparse.ArgumentParser(description='mysql client tools')
    parse.add_argument('-u','--user',dest='username',default='root',action='store_true',help='login to mysql username')
    parse.add_argument('-H','--host',dest='hostname',default='localhost',action='store_true',help='login to mysql server hostname')
    args = parse.parse_args()
    cmd = "mysql -u{} -p{}".format(args.username, args.hostname)
    print(commands.getoutput(cmd))


if __name__ == '__main__':
    main()
# 运行01
[root@yunweisn testdemo]# python3 mysql-client01.py 
ERROR 1045 (28000): Access denied for user 'root'@'localhost' (using password: YES)

0x03 pymysql模块

pymysql用于在python中连接和操作mysql数据库。

快速入门：

#!/usr/bin/python3
import pymysql

def main():
    # 创建连接
    conn = pymysql.connect(host='localhost,user='root',password='root',port=3306)
    # 创建一个游标
    cursor = conn.cursor()
    # 执行sql语句
    cursor.execute('show processlist;')
    # 获取所有的返回值
    print(cursor.fetchall())
    # 关闭
    conn.close()

    if __name__ == '__main__':
    main()

0x04 faker模块伪造数据

faker可以生成一些看起来比较真实的数据

>>> from faker import Faker
>>> faker = Faker("zh_CN")
>>> faker.name()
'莫玉珍'
>>> faker.job()
'人事经理'
>>> faker.address()
'甘肃省郑州市海港兰路E座 158840'

0x05 简单实现一下插入数据

#！/usr/bin/python3
#conding=utf-8
import argparse
import pymysql
from faker import Faker 
import sys

def args_list():
    parse = argparse.ArgumentParser('connect to mysql tools')
    parse.add_argument("--host",dest='hostname',action='store',default='localhost',help='mysql server hostname.')
    parse.add_argument("-u","--user",dest='username',action='store',default='root',help='connect to mysql server user.')
    parse.add_argument("-p","--password",dest='password',action='store',help='connect to mysql server password.')
    parse.add_argument("-P","--port",dest='port',action='store',default=3306,type=int,help='mysql server port.')
    parse.add_argument("-n","--linenumber",dest='linenumber',default=5,type=int,help='insert data line')

    args = parse.parse_args()
    return args


def conn_db(**kwargs):
    try:
        # 连接数据库
        conn = pymysql.connect(**kwargs)
        return conn # 返回conn
    except Exception as err:
        print("连接数据库错误,err=",err)
        sys.exit(1)
def create_table(cursor):
    sql = """
        create table if not exists `test`.`test001` (id int(10) not null auto_increment,
        name varchar(255) not null,
        job varchar(255) not null,
        address varchar(255) not null,
        primary key (`id`))
    """
    try:
        cursor.execute(sql)
    except Exception as err:
        print("创建表test001错误,err=",err)
        sys.exit(2)

def insert_data(cursor,linenumber,faker):
    for i in range(linenumber): # 根据用户输入的行数插入数据
        name = faker.name() # 生成名字
        job = faker.job()  # 生成职业
        address = faker.address() # 生成地址
        # 拼接sql语句
        sql = """
            insert into `test`.`test001`(name,job,address)
            values('{0}','{1}','{2}')
        """.format(name,job,address)
        print(sql)
        # 执行sql
        cursor.execute(sql)
def main():
    faker = Faker("zh_CN") # 实例化faker
    args = args_list() # 获取用户输入的信息
    #连接数据库 --> conn_db
    conn = conn_db(**dict(host=args.hostname,user=args.username,password=args.password,port=args.port))
    # 创建游标
    cursor = conn.cursor()
    # 创建表 --> create_table 
    create_table(cursor)
    # 插入数据
    insert_data(cursor,args.linenumber,faker)


if __name__ == '__main__':
    main()



#运行01
[root@yunweisn testdemo]# python3 pymysql-test02.py --host ssh.51yunwei.top -utestuser -pTestUser@123 -n 5

            insert into `test`.`test001`(name,job,address)
            values('陈秀兰','其他','河北省柳州市闵行东莞街H座 663475')


            insert into `test`.`test001`(name,job,address)
            values('陈成','学徒工','海南省上海市华龙农路H座 901586')


            insert into `test`.`test001`(name,job,address)
            values('金金凤','药品生产/质量管理','河南省莹县清浦崔街H座 341881')


            insert into `test`.`test001`(name,job,address)
            values('刘斌','汽车修理工','台湾省桂花县涪城赵路e座 673426')


            insert into `test`.`test001`(name,job,address)
            values('雷超','射频工程师','吉林省瑞市高港南宁路I座 464175')


# 登陆到数据库里看看
MariaDB [test]> select * from test001;
+----+-----------+---------------------------+----------------------------------------------+
| id | name      | job                       | address                                      |
+----+-----------+---------------------------+----------------------------------------------+
|  1 | 陈秀兰    | 其他                      | 河北省柳州市闵行东莞街H座 663475             |
|  2 | 陈成      | 学徒工                    | 海南省上海市华龙农路H座 901586               |
|  3 | 金金凤    | 药品生产/质量管理         | 河南省莹县清浦崔街H座 341881                 |
|  4 | 刘斌      | 汽车修理工                | 台湾省桂花县涪城赵路e座 673426               |
|  5 | 雷超      | 射频工程师                | 吉林省瑞市高港南宁路I座 464175               |
+----+-----------+---------------------------+----------------------------------------------+
5 rows in set (0.00 sec)

0x06 加上多进程，小工具出炉！

#！/usr/bin/python3
#conding=utf-8
import argparse
import pymysql
from faker import Faker 
import sys
import threading

def args_list():
    parse = argparse.ArgumentParser('connect to mysql tools')
    parse.add_argument("--host",dest='hostname',action='store',default='localhost',help='mysql server hostname. default localhost')
    parse.add_argument("-u","--user",dest='username',action='store',default='root',help='connect to mysql server user. default root')
    parse.add_argument("-p","--password",dest='password',action='store',help='connect to mysql server password.')
    parse.add_argument("-P","--port",dest='port',action='store',default=3306,type=int,help='mysql server port. default 3306')
    parse.add_argument("-n","--linenumber",dest='linenumber',default=500,type=int,help='insert data line , default 500')
   # 添加一个线程数量的选项
    parse.add_argument("-t","--thread_size",dest='thread_size',default=5,type=int,help='thread number，default 5')

    args = parse.parse_args()
    return args


def conn_db(**kwargs):
    try:
        conn = pymysql.connect(**kwargs)
        return conn
    except Exception as err:
        print("连接数据库错误,err=",err)
        sys.exit(1)


def create_table(cursor):
    sql = """
        create table if not exists `test`.`test001` (id int(10) not null auto_increment,
        name varchar(255) not null,
        job varchar(255) not null,
        address varchar(255) not null,
        primary key (`id`)) character set = utf8
    """
    try:
        cursor.execute(sql)
    except Exception as err:
        print("创建表test001错误,err=",err)
        sys.exit(2)

def insert_data(conn_dict,linenumber,faker):
    # 为每个线程都创建一个conn，避免资源争用
    conn = conn_db(**conn_dict)
    cursor = conn.cursor()
    for i in range(linenumber):
        name = faker.name()
        job = faker.job()
        address = faker.address()
        sql = """
            insert into `test`.`test001`(name,job,address)
            values('{0}','{1}','{2}')
        """.format(name,job,address)
        print(sql)
        cursor.execute(sql)
    cursor.execute('commit')
    # 线程操作完成后提交并关闭链接
    cursor.close()
    conn.close()
def main():
    faker = Faker("zh_CN")
    args = args_list()
    conn_dict = dict(host=args.hostname,user=args.username,password=args.password,port=args.port)
    conn = conn_db(**dict(host=args.hostname,user=args.username,password=args.password,port=args.port))
    cursor = conn.cursor()
    create_table(cursor)
    # 创建表后关闭链接
    cursor.close()
    conn.close()

    threads = [] 
    for i in range(args.thread_size):
        t = threading.Thread(target=insert_data,args=(conn_dict, args.linenumber,faker))
        threads.append(t)
        t.start()
    for t in threads:# 进程阻塞，避免未执行完成主进程退出
        t.join() 


if __name__ == '__main__':
    main()

# 运行01
[root@yunweisn testdemo]# python3 pymysql-test02.py --host localhost -uroot -proot -t 5 -n 500
......
# 查看结果
MariaDB [test]> select count(*) from test001;
+----------+
| count(*) |
+----------+
|     2530 |
+----------+
1 row in set (0.00 sec)

本文使用 文章同步助手 同步