pysqs

pysqs.py

#!/usr/bin/env python

# * simple RESTful http queue service *
# written by c4pt0r <huangdongxu1987@gmail.com>
# date 2012-8-29

import os
import tornado.ioloop
import tornado.web
import tornado.gen
import json
import redis

import time

options = None
_redis = None

def get_queue_item_block(*args, **kw):
    key = args[0]
    obj = _redis.lpop(key)
    if obj != None:
        kw['callback'](obj)
    else:
        # ugly, hope & pray
        tornado.ioloop.IOLoop.instance().add_timeout(time.time() + 0.00001, 
                                                     lambda: get_queue_item_block(*args, **kw))

class StatusHandler(tornado.web.RequestHandler):
    def get(self, queue_id):
        key = options.prefix + "queue_" + queue_id
        count = _redis.llen(key)
        info = _redis.info()
        info['queue'] = queue_id
        info['current_queue_size'] = count
        self.write(json.dumps(info))


class QueueHandler(tornado.web.RequestHandler):
    @tornado.web.asynchronous
    @tornado.gen.engine
    def get(self, queue_id):
        key = options.prefix + "queue_" + queue_id
        obj = _redis.lpop(key)
        is_block = self.get_argument('block', '0')
        if is_block == '0':
            if obj == None:
                self.write('EMPTY')
            else:
                self.write(obj)
            self.finish()
        else:
            response = yield tornado.gen.Task(get_queue_item_block, key)
            self.write(response)
            self.finish()

    def post(self, queue_id):
        key = options.prefix + "queue_" + queue_id
        value = self.get_argument('value', '')
        if len(value) > 0:
            self.write(str(_redis.rpush(key, value)))
        else:
            self.write('FAIL')

    def put(self, queue_id):
        self.post(queue_id)

# command line parser
def cmd_parse():
    from optparse import OptionParser
    parser = OptionParser()  
    parser.add_option("-r", "--redis", dest="redis", action="store", default="127.0.0.1:6379",
                      help="set redis")
    parser.add_option("-p", "--port", dest="port", action="store", default="8081",
                      help="port")  
    parser.add_option("-x", "--prefix",  
                      action="store", dest="prefix", default="pysqs_",  
                      help="key prefix of redis")  
    global options
    (options, _) = parser.parse_args()


application = tornado.web.Application([
    (r"/q/(.*)", QueueHandler),
    (r"/s/(.*)", StatusHandler),
])

if __name__ == "__main__":
    cmd_parse()
    application.listen(options.port)
    _redis = redis.Redis(options.redis.split(':')[0], 
                         int(options.redis.split(':')[1]))
    tornado.ioloop.IOLoop.instance().start()