messagequeue.py

'''
A persistent message queuing module that can be used by concurrent, distributed processes.
Uses RDBMS (e.g. MySQL) to implement the distributed, persistent, atomic and
concurrent experience. ;-)
This queue is inspired by Amazon Simple Queue Service (SQS), where any number of actors can send messages to a queue or read messages from a queue.
Fault-tolerance is implemented by locking messages read by an actor.  If an actor successfully processes a message it should delete it.
If an actor dies mysteriously, the read lock will timeout eventually and another actor will be able to read the message.

Design Goals:
    no dependency on a specific RDBMS.  
    atomic, process-level concurrency-safe.

Warning: 
    no thread-safety guarantees.
Inspiration: 
    Amazon SQS.

Usage:
Wrap a connection or a factory function in a context manager like util.NoopCM or
util.ClosingFactoryCM.
Create a MessageQueue instance or create and cache it in the module using setQueue().
Put messages on the queue.  Pop them off, either one at a time or with the generator function popGen().

The following code illustrates creating a queue which gets a new connection for each operation by using the ClosingFactory context manager.
It illustrates the prescribed idioms for reading messages, either in a for loop or with statement.
This way the state of the message queue is gracefully managed in the face of success or exceptions.

python -c'
import messagequeue, util, config;
q = messagequeue.MessageQueue("test", util.ClosingFactoryCM(config.openDbConn))
[q.send(i*i) for i in range(10)]
for m in q.readAll():
    print repr(m)
with q.read(default="empty") as m:
    print repr(m)
q.send(None)
with q.read() as m:
    print repr(m)
with q.read() as m:
    print repr(m)
'
'''


import contextlib

import dbutil


DEFAULT_LOCK_TIMEOUT = 24 * 60 * 60 # 1 day


class EmptyQueueError(Exception):
    '''
    Exception used to indicate that a queue currently has no messages.
    '''
    pass


class MessageQueue(object):
    def __init__(self, queue, manager, timeout=DEFAULT_LOCK_TIMEOUT, drop=False, create=False):
        '''
        queue: name of queue from which to send and read messages
        manager: context manager yielding a Connection.
          Typical managers are util.NoopCM(conn) to reuse a connection or
          util.ClosingFactoryCM(getConnFunc) to use a new connection each time.
        timeout: It is important that timeout be longer than it takes a
          consumer to process a message and delete it or else the message might
          be read and processed twice.
          Used to recover a message when a consumer dies while processing a
          message and before deleting it.
        '''
        self.manager = manager
        self.queue = queue
        self.timeout = timeout
        if drop:
            self._drop()
        if create:
            self._create()

    def _create(self):
        sql = '''CREATE TABLE IF NOT EXISTS message_queue ( 
                 id INT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY, 
                 queue varchar(200) NOT NULL, 
                 message blob,
                 create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                 read_time TIMESTAMP,
                 lock_time TIMESTAMP,
                 timeout INT NOT NULL,
                 locked BOOLEAN NOT NULL DEFAULT FALSE,
                 INDEX queue_index (queue) 
                 ) ENGINE = InnoDB '''
        with self.manager as conn:
            with dbutil.doTransaction(conn):
                dbutil.executeSQL(conn, sql)

    def _drop(self):
        with self.manager as conn:
            with dbutil.doTransaction(conn):
                dbutil.executeSQL(conn, 'drop table if exists message_queue')

    def send(self, message, timeout=None):
        '''
        timeout: if None, the default read lock timeout is used.  if not None, this is the number of seconds before the read lock on this message expires.
        '''
        if timeout is None:
            timeout = self.timeout
        sql = 'INSERT INTO message_queue (queue, message, timeout) VALUES (%s, %s, %s)'
        with self.manager as conn:
            with dbutil.doTransaction(conn):
                return dbutil.insertSQL(conn, sql, args=[self.queue, message, timeout])

    @contextlib.contextmanager
    def read(self, **keywords):
        '''
        Context Manager for use in with statement.  Yields a message from queue.  Message is automatically deleted from queue when with block exits.
        If an exception occurs, the message read-lock timeout is set to 0. (i.e. message is instantly available for reading.)
        default: keyword only argument.  if no messages on queue, default returned instead of EmptyQueueError exception being raised.
        '''
        try:
            i, m = self._readUnhandled()
        except EmptyQueueError:
            if keywords.has_key('default'):
                yield keywords['default']
            else:
                raise
        else:
            with self._handled(i, m) as m:
                yield m

    def readAll(self):
        '''
        Generator for use in for statement.  Yields messages from queue.  Messages are automagically deleted from queue when the yield returns.
        If an exception occurs, the last message read-lock timeout is set to 0. (i.e. message is instantly available for reading.)
        '''
        for i, m in self._readAllUnhandled():
            with self._handled(i, m) as m:
                yield m

    @contextlib.contextmanager
    def _handled(self, id, message):
        try:
            yield message
        except:
            print 'here'
            self.changeTimeout(id, 0)
            raise
        else:
            self.delete(id)
    
    def _readUnhandled(self):
        '''
        Reads the next message from the queue.
        Returns: message_id, message.
        Use message_id to delete() the message when done or to changeTimeout() of the message if necessary.
        '''
        with self.manager as conn:
            with dbutil.doTransaction(conn):
                # read first available message (pending or lock timeout)
                sql = 'SELECT id, message FROM message_queue WHERE queue = %s AND (NOT locked OR  lock_time < CURRENT_TIMESTAMP)'
                sql += ' ORDER BY id ASC LIMIT 1 FOR UPDATE '
                results = dbutil.selectSQL(conn, sql, args=[self.queue])
                if results:
                    id, message = results[0]
                    # mark message unavailable for reading for timeout seconds.
                    sql = 'UPDATE message_queue SET locked = TRUE, read_time = CURRENT_TIMESTAMP, lock_time = ADDTIME(CURRENT_TIMESTAMP, SEC_TO_TIME(timeout)) WHERE id = %s'
                    dbutil.executeSQL(conn, sql, args=[id])
                    return id, message
                else:
                    raise EmptyQueueError(str(self.queue))

    def _readAllUnhandled(self):
        while 1:
            try:
                yield self._readUnhandled()
            except EmptyQueueError:
                break

    def delete(self, id):
        sql = 'DELETE FROM message_queue WHERE id = %s '
        with self.manager as conn:
            with dbutil.doTransaction(conn):
                return dbutil.executeSQL(conn, sql, args=[id])

    def changeTimeout(self, id, timeout):
        ''' changes read lock to <timeout> seconds from now. '''
        sql = 'UPDATE message_queue SET lock_time = ADDTIME(CURRENT_TIMESTAMP, SEC_TO_TIME(%s)) WHERE id = %s'
        with self.manager as conn:
            with dbutil.doTransaction(conn):
                return dbutil.executeSQL(conn, sql, args=[timeout, id])


# last line