db.py

# -*-coding:utf-8 -*- #
# ---------------------------------------------------------------------------
# ProjectName:   ApiTest
# FileName:      db.py
# Author:       liu
# Datetime:     2024/9/7 下午3:49
# Description: 连接数据库
# ---------------------------------------------------------------------------
import pymysql

from common.read_basic_ini import Read_basic_ini
from report import log


class DB:
    def __init__(self):
        read_ini = Read_basic_ini()
        host = read_ini.get_sql('host')
        port = read_ini.get_sql('port')
        user = read_ini.get_sql('user')
        password = read_ini.get_sql('password')
        database = read_ini.get_sql('database')
        # 连接数据库
        try:
            self.con = pymysql.connect(
                host=host, user=user, password=password, database=database, port=int(port), charset='utf8'
            )
        except Exception as e:
            log.error(f'连接错误，数据：host={host}, port={port}, user={user}, password={password}, database={database}')
            raise e
        else:
            # 游标数据
            self.cursor = self.con.cursor()

    def close(self):
        # 关闭连接
        self.cursor.close()
        self.con.close()

    def delete(self, sql):
        # 判断是否为str
        if isinstance(sql, str):
            # 判断是否以delete开头
            if sql.lower().strip().startswith('delete'):
                try:
                    self.cursor.execute(sql)
                    self.con.commit()
                    log.info(f'delete方法执行{sql}语句成功!')
                except Exception as e:
                    log.error(f'执行delete方法时，发生错误，sql语句为：{sql}，错误为:{e}')
                    raise e
            else:
                log.error(f'执行delete方法时，发生错误，sql语句为：{sql}')
                raise ValueError(f'sql语句错误')
        else:
            log.error(f'执行delete方法时，发生错误，sql语句为：{sql}')
            raise ValueError(f'sql语句错误')

    def select(self, sql, n=0):
        # 判断是否为str
        if isinstance(sql, str):
            # 判断是否以delete开头
            if sql.lower().strip().startswith('select'):
                try:
                    self.cursor.execute(sql)
                    log.info(f'select方法执行{sql}语句成功!')
                except Exception as e:
                    log.error(f'执行select方法时，发生错误，sql语句为：{sql}，错误为:{e}')
                    raise e
                else:
                    select_result = self.cursor.fetchall()
                    # 获取的结果格式为：((key,),)，因此需要res[0][0]进行取值，默认取第一个的第一个，其他取值更改n来实现
                    if isinstance(n, int) and select_result:
                        try:
                            return select_result[n][0]
                        except Exception as e:
                            log.error(f"执行select方法时，获取查询结果失败，传入的n为：{n}")
                            raise e
            else:
                log.error(f'执行select方法时，发生错误，sql语句为：{sql}')
                raise ValueError(f'sql语句错误')
        else:
            log.error(f'执行select方法时，发生错误，sql语句为：{sql}')
            raise ValueError(f'sql语句错误')

if __name__ == '__main__':
    db = DB()