/
producer.py
109 lines (87 loc) · 3.58 KB
/
producer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import csv
import logging
import json
import pika
import pika.channel
from tornado import gen
from tornado.queues import Queue
from apps.base.ioloop import IOLoop
from apps.base.pika import ConnectionManager
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class Producer(object):
_channel = None
_exchange = 'user'
_exchange_type = 'topic'
_routing_key = 'user.new'
_user_queue = Queue(maxsize=100)
def setup(self, channel: pika.channel.Channel):
"""
Invoked by pika when the channel has been opened.
The channel object is passed in so we can make use of it.
"""
logger.info('Channel opened')
self._channel = channel
self._channel.add_on_close_callback(self._on_channel_closed)
self.setup_exchange(self._exchange)
IOLoop.current().spawn_callback(gen.convert_yielded(self._user_consumer()))
def setup_exchange(self, exchange_name):
"""
Setup the exchange on RabbitMQ by invoking the Exchange.Declare RPC
command. When it is complete, the on_exchange_declareok method will
be invoked by pika.
:param str exchange_name: The name of the exchange to declare
"""
logger.info('Declaring exchange %s', exchange_name)
self._channel.exchange_declare(exchange=exchange_name, exchange_type=self._exchange_type)
def _on_channel_closed(self, channel: pika.channel.Channel, reply_code: int, reply_text: str):
"""
Invoked by pika when RabbitMQ unexpectedly closes the channel.
Channels are usually closed if you attempt to do something that
violates the protocol, such as re-declare an exchange or queue with
different parameters. In this case, we'll close the connection
to shutdown the object.
:param pika.channel.Channel: The closed channel
:param int reply_code: The numeric reason the channel was closed
:param str reply_text: The text reason the channel was closed
"""
logger.warning('Channel %i was closed: (%s) %s', channel, reply_code, reply_text)
async def send(self, data):
"""
Send request action to the badges server
"""
logger.info('Sent to the user queue: {}'.format(data))
self._channel.basic_publish(
exchange=self._exchange,
routing_key=self._routing_key,
body=json.dumps(data),
properties=pika.BasicProperties(content_type='application/json', delivery_mode=2)
)
async def _user_consumer(self):
"""
Base consumer of users from local queue to dedicated
"""
async for user in self._user_queue:
try:
await self.send(user)
finally:
self._user_queue.task_done()
async def add_user(self, user):
"""
Add user to local queue, if it is full adding will be asynchronously wait resolving of it
"""
await self._user_queue.put(user)
async def fill_queue_from_csv(self, filename: str):
"""
Read CSV file with users and add all of this list to a local queue
"""
with open(filename) as users_csv:
reader = csv.DictReader(users_csv, fieldnames=['fullname', 'email'])
for user in reader:
await self.add_user(user)
producer = Producer()
connection_manager = ConnectionManager().open_channel(producer.setup)
if __name__ == '__main__':
ioloop = IOLoop.current()
ioloop.spawn_callback(gen.convert_yielded(producer.fill_queue_from_csv('data.csv')))
ioloop.start()