MySQL 数据库连接的 Python 脚本

在接口测试的过程中,需要将响应中的数据与数据库中的数据,进行断言,所以单独写了MySQL数据库连接模块。目前根据项目需求,目前实现了查询一条或多条数据,插入数据的功能。

具体实现分2个部分:

  1. 采用YAML的形式描述数据库配置
  2. 基于pymysql模块实现数据库的连接、数据查询和数据插入

数据库配置

采用YAML描述数据库配置,保存在本地

zypPlatformTest:
    host: # 数据库服务器IP地址
    port: # 数据库服务器端口号
    user: # 用户名
    password: # 密码
    db: # 数据库名称
    charset: utf8 # 连接编码

读取数据库配置,连接数据库

定义类MySQLOperation处理数据库连接、查询和插入。在__init__()函数中读取数据库配置,并连接数据库pymysql.connect
前提是安装pymysql,安装命令如pip install pymysql

    def __init__(self, config_file, database_name):

        with io.open(config_file, 'r', encoding='utf-8') as stream:
            yaml_content = yaml.load(stream)

            host = yaml_content[database_name]['host']
            port = yaml_content[database_name]['port']
            user = yaml_content[database_name]['user']
            password = yaml_content[database_name]['password']
            db_name = yaml_content[database_name]['db']
            charset = yaml_content[database_name]['charset']

        try:
            self.connection = pymysql.connect(host=host, port=port, user=user, password=password, database=db_name,
                                          charset=charset)
        except Exception as e:
            logging.error('初始化数据连接失败:%s' % e)

读取数据

连接数据库成功后,查询返回单条数据记录,或者多条数据记录
cursor.fetchone()将返回第一条数据记录,一维元组。可多次使用cursor.fetchone(),依次取得下一条数据,直到为空。
cursor.fetchall() :将返回所有结果,也可自定义返回的数据记录的条数,二维元组。

# 返回数据库的一条记录
    def select_one_record(self, query, data=""):
        logging.info('query:%s  data:%s' % (query, data))

        with self.connection.cursor() as cursor:
            if data:
                cursor.execute(query, data)
            else:
                cursor.execute(query)
            query_result = cursor.fetchone()

        return query_result

    def select_many_record(self, query, data=""):
        logging.info('query:%s  data:%s' % (query, data))

        with self.connection.cursor() as cursor:
            if data:
                numbers = cursor.execute(query, data)
            else:
                numbers = cursor.execute(query)
            query_result = cursor.fetchmany(numbers)

        return query_result

插入数据

连接数据库成功后,目前只实现了插入单条数据记录。

def execute_insert(self, query, data=''):
    logging.info('query:%s  data:%s' % (query, data))
    with self.connection.cursor() as cursor:
        if data:
            cursor.execute(query, data)
        else:
            cursor.execute(query)

    # connection is not autocommit by default. So you must commit to save
    self.connection.commit()

断开数据库连接

def close_connection(self):
    self.connection.close()
关闭