forked from DxvLwRYF/apiAutoTest
67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
#!/usr/bin/env/python3
|
||
# -*- coding:utf-8 -*-
|
||
"""
|
||
@project: apiAutoTest
|
||
@author: zy7y
|
||
@file: db.py
|
||
@ide: PyCharm
|
||
@time: 2020/12/4
|
||
@desc: 数据库连接,目前只支持mysql ，且个人认为用到最多的操作应该是查询所以其他todo
|
||
"""
|
||
import json
|
||
from datetime import datetime
|
||
from typing import Union
|
||
|
||
import pymysql
|
||
|
||
from tools.read_file import ReadFile
|
||
|
||
|
||
class DB:
|
||
mysql = ReadFile.read_config('$.database')
|
||
|
||
def __init__(self):
|
||
"""
|
||
初始化数据库连接，并指定查询的结果集以字典形式返回
|
||
"""
|
||
self.connection = pymysql.connect(
|
||
host=self.mysql['host'],
|
||
port=self.mysql['port'],
|
||
user=self.mysql['user'],
|
||
password=self.mysql['password'],
|
||
db=self.mysql['db_name'],
|
||
charset=self.mysql.get('charset', 'utf8mb4'),
|
||
cursorclass=pymysql.cursors.DictCursor
|
||
)
|
||
|
||
def execute_sql(self, sql: str) -> Union[dict, None]:
|
||
"""
|
||
执行sql语句方法，查询所有结果的sql只会返回一条结果（
|
||
比如说： 使用select * from cases , 结果将只会返回第一条数据 {'id': 1, 'name': 'updatehahaha', 'path': None, 'body': None, 'expected': '{"msg": "你好"}', 'api_id': 1, 'create_at': '2021-05-17 17:23:54', 'update_at': '2021-05-17 17:23:54'}
|
||
|
||
），支持select， delete， insert， update
|
||
:param sql: sql语句
|
||
:return: select 语句 如果有结果则会返回 对应结果字典，delete，insert，update 将返回None
|
||
"""
|
||
with self.connection.cursor() as cursor:
|
||
cursor.execute(sql)
|
||
result = cursor.fetchone()
|
||
# 使用commit解决查询数据出现概率查错问题
|
||
self.connection.commit()
|
||
return self.verify(result)
|
||
|
||
def verify(self, result: dict) -> Union[dict, None]:
|
||
"""验证结果能否被json.dumps序列化"""
|
||
# 尝试变成字符串，解决datetime 无法被json 序列化问题
|
||
try:
|
||
json.dumps(result)
|
||
except TypeError: # TypeError: Object of type datetime is not JSON serializable
|
||
for k, v in result.items():
|
||
if isinstance(v, datetime):
|
||
result[k] = str(v)
|
||
return result
|
||
|
||
def close(self):
|
||
"""关闭数据库连接"""
|
||
self.connection.close()
|