import MySQLdb
import time
import sys
import logging

class MySQLdbSession:

    def __init__(self, mysqldb_connect_params, max_reconnect_attempts=5, pause_between_attempts=1, is_ss=False):
        self.mysqldb_connect_params = mysqldb_connect_params
        self.reconnect_attempts = 0
        self.wait_time = pause_between_attempts #in seconds
        self.max_reconnect_attempts = max_reconnect_attempts
        if is_ss:
            raise NotImplemented
        self._refresh_cursor()

    def _refresh_cursor(self):
        self.db_conn = MySQLdb.connect(**self.mysqldb_connect_params)
        self.cursor = self.db_conn.cursor()

    def _try_func(self, func, args):
        '''
        :param func: function - the function to try
        :param args: tuple - args to be tried with the function
        '''
        try:
            retval = func(*args)
        except Exception as e:
            logging.warning("There was an exception...")
            logging.warning(sys.exc_info())
            is_server_restart_exception = e.message == "MySQL server has gone away"
            too_many_reconnects = self.reconnect_attempts >= self.max_reconnect_attempts
            if is_server_restart_exception and not too_many_reconnects:
                self.reconnect_attempts += 1
                logging.info("Sleeping...")
                time.sleep(self.wait_time)
                logging.info("Trying to reconnect...")
                self._refresh_cursor()
                retval = self._try_func(func, args)
            else:
                raise e
        self.reconnect_attempts = 0
        return retval

    def execute(self, *args):
        return self._try_func(self.cursor.execute, args)

    def executemany(self, *args):
        return self._try_func(self.cursor.executemany, args)

    def fetchall(self, *args):
        return self._try_func(self.cursor.fetchall, args)

    def commit(self, *args):
        return self._try_func(self.db_conn.commit, args)