import pymysql
import time
import pandas as pd

class DBDownload:

    def __init__(self, host='', port='', db='', user='', password='', mode=''):
        self.host = host
        self.port = port
        self.db = db
        self.user = user
        self.password = password
        self.mode=mode
        pass

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def connect(self):
        conn_success_flag = 0
        while not conn_success_flag:
            try:
                self.conn = pymysql.connect(host=self.host, port=self.port, user=self.user, password=self.password, database=self.db)
            except Exception as e:
                conn_success_flag = 0
                print("数据库连接失败 :{}".format(e))
                time.sleep(5)
            else:
                conn_success_flag = 1
                print('数据库连接成功！')
                self.cursor = self.conn.cursor()

    def getdata(self,*param,tablename,sn,timename,st,sp,factory):
        print('数据获取中......')
        param=list(param)
        str=''
        for i in range(len(param)):
            if i<1:
                str=str+param[i]
            else:
                str=str+','+param[i]
        if self.mode==1:
            self.cursor.execute("select %s from %s where sn='%s' order by id desc limit 1" %(str,tablename,sn))
        elif self.mode==0:
<<<<<<< HEAD
            if len(sn)>1:
                self.cursor.execute("select %s from %s where sn='%s' and %s between '%s' and '%s'" %(str,tablename,sn,timename,st,sp))
            else:
                self.cursor.execute("select %s from %s where %s between '%s' and '%s'" %(str,tablename,timename,st,sp))
=======
            self.cursor.execute("select %s from %s where sn='%s' and %s between '%s' and '%s'" %(str,tablename,sn,timename,st,sp))
>>>>>>> master
        elif self.mode==2:
            self.cursor.execute("select %s from %s where factory='%s'" %(str,tablename,factory))
        elif self.mode==3:
            self.cursor.execute("select %s from %s" %(str,tablename))
        elif self.mode==4:
            self.cursor.execute("select %s from %s where qrcode='%s' order by id desc limit 1" %(str,tablename,sn))
<<<<<<< HEAD
        elif self.mode==5:
            self.cursor.execute("select %s from %s where info='%s'" %(str,tablename,timename))
=======
>>>>>>> master
        res = self.cursor.fetchall()
        df_res = pd.DataFrame(res, columns=param)
        df_res = df_res.reset_index(drop=True)
        return(df_res)

    def close(self):
        try:
            self.cursor.close()
            self.conn.close()
        except Exception as e:
            print(e)
        else:
            print('数据库已断开连接！')