#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
import pymysql
import psycopg2


class DB():

    def __init__(self, DB_HOST, DB_PORT, DB_USER, DB_PWD, DB_NAME, DB='mysql'):
        self.DB_HOST = DB_HOST  # 设置MYSQL地址
        self.DB_PORT = DB_PORT  # 设置端口号
        self.DB_USER = DB_USER  # 设置用户名
        self.DB_PWD = DB_PWD  # 设置密码
        self.DB_NAME = DB_NAME  # 数据库名
        self.DB = DB  # 连接哪种数据库，mysql或者pgsql
        self.conn = self.getConnection()

    def getConnection(self):
        # 连接哪种数据库
        if self.DB == 'pgsql':
            return psycopg2.connect(
                host=self.DB_HOST,
                port=self.DB_PORT,
                user=self.DB_USER,
                password=self.DB_PWD,
                database=self.DB_NAME,

            )
        elif self.DB == 'mysql':
            return pymysql.Connect(
                host=self.DB_HOST,
                port=self.DB_PORT,
                user=self.DB_USER,
                password=self.DB_PWD,
                database=self.DB_NAME,
                charset='utf8'
            )
        else:
            raise ("请选择正确连接数据库方式，mysql或者pgsql")
    def close(self):
        self._close(self.conn.cursor())
    # 查询方法
    def query(self, sqlString):
        cursor = self.conn.cursor()
        cursor.execute(sqlString)
        returnData = cursor.fetchall()
        # self._close(cursor)
        return returnData


    def getConn(self):
        return self.conn

    def getCursor(self):
        return self.conn.cursor()

    # 获取单条数据
    def find_one(self, sqlString):
        cursor = self.conn.cursor()
        cursor.execute(sqlString)
        returnData = cursor.fetchone()
        self._close(cursor)
        return returnData

    # 不带参数的更新方法
    def update(self, sqlString):
        cursor = self.conn.cursor()
        cursor.execute(sqlString)
        self.conn.commit()

    # 带参数的更新方法
    def updateByParam(self, sqlString, params):
        cursor = self.conn.cursor()
        cursor.execute(sqlString, params)
        self.conn.commit()
        self._close(cursor)

    # 关闭回收资源
    def _close(self, cursor):
        cursor.close()
        self.conn.close()


# if __name__ == "__main__":
    # db = DB('localhost', 3306, 'root', '123456', 'zuqiu')
    # sql = "show create table dc_sclassinfo;"
    # result = db.query(sql)
    # pycomm.toLog(result)