root/aws/aws/sqs.py

from StringIO import StringIO
import urllib
import mimetypes
import os, os.path
import cPickle as pickle

from twisted.internet import task, reactor, defer, protocol
from twisted.application import service
from twisted.python import usage

from twisted.web2.client import http
from twisted.web2 import stream, http_headers

from lxml import etree

from aws.client import rest 
from aws import s3_example, util

AMAZON_SQS_NS = "http://queue.amazonaws.com/doc/2007-05-01/"
default_ns = {"sqs": AMAZON_SQS_NS}

class Options(usage.Options):
    optParameters = [
            ['access_key', "k", 'I SHOULD BE A KEY'],
            ['secret_access_key', "s", 'I SHOULD BE THE ULTRA SUPER SECRET KEY'],
            ['base_url', None, "http://queue.amazonaws.com"],
            ['aws_version', None, "2007-05-01"],
        ]
    
def _print(s):
    print s
    return s

class SQS(object):
    def __init__(self, config):
        self.config = config
        self.access_key = config['access_key']
        self.secret_access_key = config['secret_access_key']
        self.base_url = config['base_url']
    
    # interface
    def postQueue(self, name):
        options = {"QueueName": name}
        query = urllib.urlencode(options)

        url = self.base_url + "/?" + query
        
        params = dict(url=url, method="POST")
    

        def _checkErrors(r):
            #print r.headers.getAllRawHeaders()
            r.stream.read().addCallback(_print)
            if r.code != 200:
                raise Exception("bad status", r)
            return r

        def _grabTheUrl(r):
            return r.xpath("//sqs:QueueUrl/text()", namespaces=default_ns)[0]

        d = self.getPage(**params)
        d.addCallback(_checkErrors)
        d.addCallback(build_xml_response)
        d.addCallback(_grabTheUrl)
        d.addBoth(_print)

        return d
    def listQueues(self, prefix=None):
        url = self.base_url
        if prefix:
            query = urllib.urlencode({"QueueNamePrefix": prefix})
            url = url + "/?" + query

        params = dict(url=url, method="GET")

        def _checkErrors(r):
            #r.stream.read().addCallback(_print)
            if r.code != 200:
                raise Exception("bad status", r)
            return r

        def _grabTheUrls(r):
            return r.xpath("//sqs:QueueUrl/text()", namespaces=default_ns)

        d = self.getPage(**params)
        d.addCallback(_checkErrors)
        d.addCallback(build_xml_response)
        d.addCallback(_grabTheUrls)
        return d
    def getQueue(self, queue, max=None, timeout=None):
        url = queue + "/front"

        options = {}
        if max is not None:
            options['NumberOfMessages'] = max
        if timeout is not None:
            options['VisibilityTimeout'] = timeout
        if options:
            query = urllib.urlencode(options)
            url = url + "?" + query
        params = dict(url=url, method="GET")

        def _checkErrors(r):
            r.stream.read().addCallback(_print)
            if r.code != 200:
                raise Exception("bad status", r)
            return r

        def _grabTheMessages(r):
            elms = r.xpath("//sqs:Message", namespaces=default_ns)
            messages = []
            for elm in elms:
                messages.append(SQSMessage(
                    elm.xpath("sqs:MessageId/text()", 
                              namespaces=default_ns)[0],
                    elm.xpath("sqs:MessageBody/text()", 
                              namespaces=default_ns)[0],
                    queue))
                    
            return messages

        d = self.getPage(**params)
        d.addCallback(_checkErrors)
        d.addCallback(build_xml_response)
        d.addCallback(_grabTheMessages)
        return d
    def putMessage(self, queue, data, headers=None):
        url = "%s/back" % (queue)
        if not headers:
            headers = {}
        
        headers.setdefault("Content-type", "text/plain")
                   
        params = dict(url=url, method="PUT", headers=headers, data=data)

        def _checkErrors(r):
            print r.headers.getAllRawHeaders()
            r.stream.read().addCallback(_print)
            if r.code != 200:
                raise Exception("bad status", r)
            return r

        d = self.getPage(**params)
        d.addCallback(_checkErrors)
        return d
    def putMessageFromFile(self, queue, file, headers=None):
        url = "%s/back" % (queue)
        if not headers:
            headers = {}
        
        content_type, content_encoding = mimetypes.guess_type(file.name)
        headers.setdefault("Content-type", content_type)
        #headers.setdefault("Content-encoding", content_encoding)
        
        data = stream.FileStream(file)

        return self.putMessage(queue=queue, data=data, headers=headers)
    def getMessage(self, queue, message_id):
        url = queue + "/" + message_id

        params = dict(url=url, method="GET")

        def _checkErrors(r):
            if r.code != 200:
                raise Exception("bad status", r)
            return r

        def _grabTheMessage(r):
            elms = r.xpath("//sqs:Message", namespaces=default_ns)
            messages = []
            for elm in elms:
                messages.append(SQSMessage(
                    elm.xpath("sqs:MessageId/text()", 
                              namespaces=default_ns)[0],
                    elm.xpath("sqs:MessageBody/text()", 
                              namespaces=default_ns)[0],
                    queue))
                    
            return messages[0]

        d = self.getPage(**params)
        d.addBoth(_checkErrors)
        d.addCallback(build_xml_response)
        d.addCallback(_grabTheMessage)
        return d
    def deleteMessage(self, queue, message_id):
        url = queue + "/" + message_id

        params = dict(url=url, method="DELETE")

        def _checkErrors(r):
            if r.code != 200:
                raise Exception("bad status", r)
            return r

        d = self.getPage(**params)
        d.addBoth(_checkErrors)
        return d
    
    # util
    def getPage(self, **kw):
        kw.setdefault("access_key", self.access_key)
        kw.setdefault("secret_access_key", self.secret_access_key)
        kw.setdefault("headers", {})
        kw['headers'].setdefault("AWS-Version", self.config['aws_version'])
        kw['headers'].setdefault("Content-type", "text/plain")
        return rest.getPage(**kw)

class SQSMessage(object):
    def __init__(self, id, body, queue):
        self.id = id
        self.body = body
        self.queue = queue
    def __str__(self):
        return "<SQSMessage id='%s' data='%s'>" % (self.id, self.body.strip())
    def __repr__(self):
        return str(self)


def build_xml_response(r):
    d = util.exhaust_stream(r.stream)
    d.addCallback(StringIO)
    d.addCallback(etree.parse)
    return d


def makeService(config):
    s = SQS(config)
    return s


### Support for Amazon S3 example code
class OldResponse(object):
    def __init__(self, status, headers, body):
        self.headers = headers
        self.msg = headers
        self.body = StringIO(body)
        self.status = status
    
    def read(self, *args, **kw):
        return self.body.read(*args, **kw)
def build_old_response(response):
    def _assemble(b):
        headers = dict(response.headers.getAllRawHeaders())
        o = OldResponse(response.code, headers, b)
        return o
    d = util.exhaust_stream(response.stream)
    d.addCallback(_assemble)
    return d