pyrabbit.py

   1 # -*- coding: utf-8 -*-
   2 ##############################################################################
   3 #
   4 #    pyrabbit, a Python library for easy playing with RabbitMQ
   5 #    Copyright (C) 2013 by Javier Sancho Fernandez <jsf at jsancho dot org>
   6 #
   7 #    This program is free software: you can redistribute it and/or modify
   8 #    it under the terms of the GNU General Public License as published by
   9 #    the Free Software Foundation, either version 3 of the License, or
  10 #    (at your option) any later version.
  11 #
  12 #    This program is distributed in the hope that it will be useful,
  13 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 #    GNU General Public License for more details.
  16 #
  17 #    You should have received a copy of the GNU General Public License
  18 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
  19 #
  20 ##############################################################################
  21
  22 import errno
  23 import os
  24 import pika
  25 import signal
  26 import uuid
  27
  28 class TimeoutError(Exception):
  29     pass
  30
  31 class Connection(object):
  32     def __init__(self, host):
  33         self.host = host
  34         self.open(host=self.host)
  35
  36     def open(self, host=None):
  37         if host is None:
  38             host = self.host
  39         self.connection = pika.BlockingConnection(pika.ConnectionParameters(host=host))
  40         self.channel = self.connection.channel()
  41         self.callback_queue = None
  42
  43     def close(self):
  44         self.channel.cancel()
  45         self.connection.close()
  46
  47     def __getattr__(self, queue_name):
  48         return Queue(self, queue_name)
  49
  50     def __getitem__(self, queue_name):
  51         return self.__getattr__(queue_name)
  52
  53     def receive(self, queues=[], timeout=0):
  54         if len(queues) == 0:
  55             return None
  56
  57         def _handle_timeout(signum, frame):
  58             raise TimeoutError(os.strerror(errno.ETIME))
  59         signal.signal(signal.SIGALRM, _handle_timeout)
  60         signal.alarm(timeout)
  61
  62         res = None
  63         try:
  64             method = None
  65             i = 0
  66             while method is None:
  67                 method, properties, body = self.channel.basic_get(queues[i])
  68                 if i == len(queues) - 1:
  69                     i = 0
  70                 else:
  71                     i += 1
  72             res = Message(self, method, properties, body)
  73         finally:
  74             signal.alarm(0)
  75
  76         return res
  77
  78
  79 class Queue(object):
  80     def __init__(self, connection, queue_name):
  81         self.connection = connection
  82         self.queue_name = queue_name
  83         self.connection.channel.queue_declare(queue=self.queue_name)
  84
  85     def receive(self, timeout=0):
  86         return self.connection.receive(queues=[self.queue_name], timeout=timeout)
  87
  88     def send(self, body, wait_response=False, timeout=0):
  89         properties = None
  90
  91         if wait_response:
  92             corr_id = str(uuid.uuid4())
  93             if self.connection.callback_queue is None:
  94                 self.connection.callback_queue = self.connection.channel.queue_declare(exclusive=True).method.queue
  95             properties = pika.BasicProperties(
  96                 reply_to=self.connection.callback_queue,
  97                 correlation_id=corr_id,
  98                 )
  99         else:
 100             properties = pika.BasicProperties()
 101
 102         self.connection.channel.basic_publish(exchange='',
 103                                               routing_key=self.queue_name,
 104                                               properties=properties,
 105                                               body=body)
 106
 107         response = None
 108         if wait_response:
 109             response = self.connection.receive(queues=[self.connection.callback_queue], timeout=timeout)
 110
 111         return response
 112
 113
 114 class Message(object):
 115     def __init__(self, connection, method, properties, body):
 116         self.connection = connection
 117         self.method = method
 118         self.properties = properties
 119         self.body = body
 120         self._ack = False
 121
 122     def ack(self):
 123         if not self._ack:
 124             self.connection.channel.basic_ack(delivery_tag=self.method.delivery_tag)
 125             self._ack = True
 126
 127     def response(self, body):
 128         if self.properties.reply_to:
 129             self.connection.channel.basic_publish(exchange='',
 130                                                   routing_key=self.properties.reply_to,
 131                                                   properties=pika.BasicProperties(correlation_id=self.properties.correlation_id),
 132                                                   body=body)