Receiving from more than one queue (using basic_get instead consume)
[pyrabbit.git] / 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)