import pymysql

def reconnect_mysql(func):
    def wrapper(self, *args, **kwargs):
        try:
            self.connection.ping(reconnect=True)
        except pymysql.OperationalError:
            self.connection = pymysql.connect(host=self.host, port=self.port, user=self.user, password=self.password,
                                              db=self.db, charset='utf8mb4')
        finally:
            return func(self, *args, **kwargs)

    return wrapper
 
class Data(object):
    def __init__(self):
        self.connection = None
        self.cursor = None

    def select_data(self, *args, **kwargs):
        pass

    def count(self, table_name, sql=None):
        if not sql:
            sql = 'select count(*) from %s' % table_name

        self.cursor.execute(sql)
        return self.cursor.fetchall()[0][0]

    def column_names(self, table_name):
        pass
      
      
class MysqlData(Data):
    def __init__(self, host='127.0.0.1', port=3306, user=None, password=None, db=None):
        super(MysqlData, self).__init__()
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.db = db
        # # 连接mysql
        self.connection = pymysql.connect(host=self.host, port=self.port, user=self.user, password=self.password,
                                          db=self.db, charset='utf8mb4')
        self.cursor = self.connection.cursor()

    @reconnect_mysql
    def insert_data(self, mysql_sql, data):
        result = self.cursor.execute(mysql_sql, data)
        self.connection.commit()

        return result

    @reconnect_mysql
    def insert_many_data(self, mysql_sql, data):
        result = self.cursor.executemany(mysql_sql, data)
        self.connection.commit()

        return result

    @reconnect_mysql
    def column_names(self, table_name):
        sql = "select DISTINCT COLUMN_NAME from information_schema.COLUMNS where table_name = '%s'" % table_name
        self.cursor.execute(sql)
        return [col[0] for col in self.cursor.fetchall()]

    @reconnect_mysql
    def select_data(self, sql):
        self.cursor.execute(sql)
        rows = self.cursor.fetchall()

        return rows

    @reconnect_mysql
    def update_data(self, sql):
        result = self.cursor.execute(sql)
        self.connection.commit()

        return result

    def close_db(self):
        self.connection.close()
        self.cursor.close()
        
        
if __name__ == '__main__':
    mysql = MysqlData(host='1.1.1.1',
                                   port=3306,
                                   user='aaa',
                                   password='bbb',
                                   db='ccc')