# Source code for pyzor.engines.mysql

    """MySQLdb database engine."""

    import time
    import logging
    import datetime
    import itertools
    import functools
    import threading

    try:
        import Queue
    except ImportError:
        import queue as Queue

    try:
        import MySQLdb
        import MySQLdb.cursors
        _has_mysql = True
    except ImportError:
        _has_mysql = False

    from pyzor.engines.common import *

    class MySQLDBHandle(BaseEngine):
        absolute_source = False
        handles_one_step = True
        # The table must already exist, and have this schema:
        #   CREATE TABLE `public` (
        #   `digest` char(40) default NULL,
        #   `r_count` int(11) default NULL,
        #   `wl_count` int(11) default NULL,
        #   `r_entered` datetime default NULL,
        #   `wl_entered` datetime default NULL,
        #   `r_updated` datetime default NULL,
        #   `wl_updated` datetime default NULL,
        #   PRIMARY KEY  (`digest`)
        #   )
        # XXX Re-organising might be faster with a r_updated index.  However,
        # XXX the re-organisation time isn't that important, and that would
        # XXX (slightly) slow down all inserts, so we leave it for now.
        reorganize_period = 3600 * 24  # 1 day
        reconnect_period = 60  # seconds
        log = logging.getLogger("pyzord")

        def __init__(self, fn, mode, max_age=None):
            self.max_age = max_age
            self.db = None
            # The 'fn' is host,user,password,db,table.  We ignore mode.
            # We store the authentication details so that we can reconnect if
            # necessary.
            self.host, self.user, self.passwd, self.db_name, \
                self.table_name = fn.split(",")
            self.last_connect_attempt = 0  # We have never connected.
            self.reorganize_timer = None
            self.reconnect()
            self.start_reorganizing()

        def _get_new_connection(self):
            """Returns a new db connection."""
            db = MySQLdb.connect(host=self.host, user=self.user,
                                 db=self.db_name, passwd=self.passwd)
            db.autocommit(True)
            return db

        def _check_reconnect_time(self):
            if time.time() - self.last_connect_attempt < self.reconnect_period:
                # Too soon to reconnect.
                self.log.debug("Can't reconnect until %s",
                               (time.ctime(self.last_connect_attempt +
                                           self.reconnect_period)))
                return False
            return True

        def reconnect(self):
            if not self._check_reconnect_time():
                return
            if self.db:
                try:
                    self.db.close()
                except MySQLdb.Error:
                    pass
            try:
                self.db = self._get_new_connection()
            except MySQLdb.Error as e:
                self.log.error("Unable to connect to database: %s", e)
                self.db = None
            # Keep track of when we connected, so that we don't retry too often.
            self.last_connect_attempt = time.time()

        def _iter(self, db):
            c = db.cursor(cursorclass=MySQLdb.cursors.SSCursor)
            c.execute("SELECT digest FROM %s" % self.table_name)
            while True:
                row = c.fetchone()
                if not row:
                    break
                yield row[0]
            c.close()

        def __iter__(self):
            return self._safe_call("iter", self._iter, ())

        def _iteritems(self, db):
            c = db.cursor(cursorclass=MySQLdb.cursors.SSCursor)
            c.execute("SELECT digest, r_count, wl_count, r_entered, r_updated, "
                      "wl_entered, wl_updated FROM %s" % self.table_name)
            while True:
                row = c.fetchone()
                if not row:
                    break
                yield row[0], Record(*row[1:])
            c.close()

        def iteritems(self):
            return self._safe_call("iteritems", self._iteritems, ())

        def items(self):
            return list(self._safe_call("iteritems", self._iteritems, ()))

        def __del__(self):
            """Close the database when the object is no longer needed."""
            try:
                if self.db:
                    self.db.close()
            except MySQLdb.Error:
                pass

        def _safe_call(self, name, method, args):
            try:
                return method(*args, db=self.db)
            except (MySQLdb.Error, AttributeError) as ex:
                self.log.error("%s failed: %s", name, ex)
                self.reconnect()
                # Retrying just complicates the logic - we don't really care if
                # a single query fails (and it's possible that it would fail)
                # on the second attempt anyway.  Any exceptions are caught by
                # the server, and a 'nice' message provided to the caller.
                raise DatabaseError("Database temporarily unavailable.")

        def report(self, keys):
            return self._safe_call("report", self._report, (keys,))

        def whitelist(self, keys):
            return self._safe_call("whitelist", self._whitelist, (keys,))

        def __getitem__(self, key):
            return self._safe_call("getitem", self._really__getitem__, (key,))

        def __setitem__(self, key, value):
            return self._safe_call("setitem", self._really__setitem__,
                                   (key, value))

        def __delitem__(self, key):
            return self._safe_call("delitem", self._really__delitem__, (key,))

        def _report(self, keys, db=None):
            c = db.cursor()
            try:
                c.executemany("INSERT INTO %s (digest, r_count, wl_count, "
                              "r_entered, r_updated, wl_entered, wl_updated) "
                              "VALUES (%%s, 1, 0, NOW(), NOW(), NOW(), NOW()) ON "
                              "DUPLICATE KEY UPDATE r_count=r_count+1, "
                              "r_updated=NOW()" % self.table_name,
                              itertools.imap(lambda key: (key,), keys))
            finally:
                c.close()

        def _whitelist(self, keys, db=None):
            c = db.cursor()
            try:
                c.executemany("INSERT INTO %s (digest, r_count, wl_count, "
                              "r_entered, r_updated, wl_entered, wl_updated) "
                              "VALUES (%%s, 0, 1, NOW(), NOW(), NOW(), NOW()) ON "
                              "DUPLICATE KEY UPDATE wl_count=wl_count+1, "
                              "wl_updated=NOW()" % self.table_name,
                              itertools.imap(lambda key: (key,), keys))
            finally:
                c.close()

        def _really__getitem__(self, key, db=None):
            """__getitem__ without the exception handling."""
            c = db.cursor()
            # The order here must match the order of the arguments to the
            # Record constructor.
            c.execute("SELECT r_count, wl_count, r_entered, r_updated, "
                      "wl_entered, wl_updated FROM %s WHERE digest=%%s" %
                      self.table_name, (key,))
            try:
                try:
                    return Record(*c.fetchone())
                except TypeError:
                    # fetchone() returned None, i.e. there is no such record
                    raise KeyError()
            finally:
                c.close()

        def _really__setitem__(self, key, value, db=None):
            """__setitem__ without the exception handling."""
            c = db.cursor()
            try:
                c.execute("INSERT INTO %s (digest, r_count, wl_count, "
                          "r_entered, r_updated, wl_entered, wl_updated) "
                          "VALUES (%%s, %%s, %%s, %%s, %%s, %%s, %%s) ON "
                          "DUPLICATE KEY UPDATE r_count=%%s, wl_count=%%s, "
                          "r_entered=%%s, r_updated=%%s, wl_entered=%%s, "
                          "wl_updated=%%s" % self.table_name,
                          (key, value.r_count, value.wl_count, value.r_entered,
                           value.r_updated, value.wl_entered, value.wl_updated,
                           value.r_count, value.wl_count, value.r_entered,
                           value.r_updated, value.wl_entered, value.wl_updated))
            finally:
                c.close()

        def _really__delitem__(self, key, db=None):
            """__delitem__ without the exception handling."""
            c = db.cursor()
            try:
                c.execute("DELETE FROM %s WHERE digest=%%s" % self.table_name,
                          (key,))
            finally:
                c.close()

        def start_reorganizing(self):
            if not self.max_age:
                return
            self.log.debug("reorganizing the database")
            breakpoint = (datetime.datetime.now() -
                          datetime.timedelta(seconds=self.max_age))
            db = self._get_new_connection()
            c = db.cursor()
            try:
                c.execute("DELETE FROM %s WHERE r_updated<%%s" %
                          self.table_name, (breakpoint,))
            except (MySQLdb.Error, AttributeError) as e:
                self.log.warn("Unable to reorganise: %s", e)
            finally:
                c.close()
                db.close()
            self.reorganize_timer = threading.Timer(self.reorganize_period,
                                                    self.start_reorganizing)
            self.reorganize_timer.setDaemon(True)
            self.reorganize_timer.start()

        @classmethod
        def get_prefork_connections(cls, fn, mode, max_age=None):
            """Yields a number of database connections suitable for a Pyzor
            pre-fork server.
            """
            # Only run the reorganize timer in the first child process.
            yield functools.partial(cls, fn, mode, max_age=max_age)
            while True:
                yield functools.partial(cls, fn, mode, max_age=None)

    class ThreadedMySQLDBHandle(MySQLDBHandle):
        def __init__(self, fn, mode, max_age=None, bound=None):
            self.bound = bound
            if self.bound:
                self.db_queue = Queue.Queue()
            MySQLDBHandle.__init__(self, fn, mode, max_age=max_age)

        def _get_connection(self):
            if self.bound:
                return self.db_queue.get()
            else:
                return self._get_new_connection()

        def _release_connection(self, db):
            if self.bound:
                self.db_queue.put(db)
            else:
                db.close()

        def _safe_call(self, name, method, args):
            db = self._get_connection()
            try:
                return method(*args, db=db)
            except (MySQLdb.Error, AttributeError) as ex:
                self.log.error("%s failed: %s", name, ex)
                if not self.bound:
                    raise DatabaseError("Database temporarily unavailable.")
                try:
                    # Connection might be timeout, ping and retry
                    db.ping(True)
                    return method(*args, db=db)
                except (MySQLdb.Error, AttributeError) as ex:
                    # attempt a new connection, if we can retry
                    db = self._reconnect(db)
                    raise DatabaseError("Database temporarily unavailable.")
            finally:
                self._release_connection(db)

        def reconnect(self):
            if not self.bound:
                return
            for _ in xrange(self.bound):
                self.db_queue.put(self._get_new_connection())

        def _reconnect(self, db):
            if not self._check_reconnect_time():
                return db
            else:
                self.last_connect_attempt = time.time()
                return self._get_new_connection()

        def __del__(self):
            if not self.bound:
                return
            for db in iter(self.db_queue.get_nowait):
                try:
                    db.close()
                except MySQLdb.Error:
                    continue
                except Queue.Empty:
                    break

    class ProcessMySQLDBHandle(MySQLDBHandle):
        def __init__(self, fn, mode, max_age=None):
            MySQLDBHandle.__init__(self, fn, mode, max_age=max_age)

        def reconnect(self):
            pass

        def __del__(self):
            pass

        def _safe_call(self, name, method, args):
            db = None
            try:
                db = self._get_new_connection()
                return method(*args, db=db)
            except (MySQLdb.Error, AttributeError) as ex:
                self.log.error("%s failed: %s", name, ex)
                raise DatabaseError("Database temporarily unavailable.")
            finally:
                if db is not None:
                    db.close()

    if not _has_mysql:
        handle = DBHandle(single_threaded=None,
                          multi_threaded=None,
                          multi_processing=None,
                          prefork=None)
    else:
        handle = DBHandle(single_threaded=MySQLDBHandle,
                          multi_threaded=ThreadedMySQLDBHandle,
                          multi_processing=ProcessMySQLDBHandle,
                          prefork=MySQLDBHandle)