def migrate_tasks(source, dest, timeout=1.0, app=None,
migrate=None, callback=None):
state = State()
app = app_or_default(app)
def update_state(body, message):
state.count += 1
producer = app.amqp.TaskPublisher(dest)
if migrate is None:
migrate = partial(migrate_task, producer)
if callback is not None:
callback = partial(callback, state)
consumer = app.amqp.get_task_consumer(source)
consumer.register_callback(update_state)
consumer.register_callback(callback)
consumer.register_callback(migrate)
# declare all queues on the new broker.
for queue in consumer.queues:
queue(producer.channel).declare()
try:
_, mcount, _ = queue(consumer.channel).queue_declare(passive=True)
if mcount:
state.total_apx += mcount
except source.channel_errors + (StdChannelError, ):
pass
# start migrating messages.
with consumer:
try:
for _ in eventloop(source, timeout=timeout):
pass
except socket.timeout:
return
def consume(self):
"""Consumes all messages on the queues
"""
self.callbacks.append(self.ack_message)
c = messaging.Consumer(self.connection, queues=self.queues, callbacks=self.callbacks, auto_declare=False)
c.qos(prefetch_count=self.prefetch_count)
with c:
for _ in eventloop(self.connection, timeout=10, ignore_timeouts=True):
pass
def consume():
"""Consume jobs from queue and do work."""
try:
with Connection('amqp://*****:*****@localhost:5672/%2F') as conn:
conn.ensure_connection()
exchange = Exchange('data_staged', type='direct')
queue = Queue('data_staged', exchange, routing_key='data_staged')
channel = conn.channel()
channel.basic_qos(0, 1, False)
logger.info("Starting worker.")
with conn.Consumer(queue, channel, callbacks=[callback]) as consumer:
for _ in eventloop(conn, timeout=1, ignore_timeouts=True):
pass
except KeyboardInterrupt:
logger.info("Worker was shut down.")
def drain_consumer(consumer, limit=1, timeout=None, callbacks=None,
ignore_timeouts=True):
acc = deque()
def on_message(body, message):
acc.append((body, message))
consumer.callbacks = [on_message] + (callbacks or [])
with consumer:
client = consumer.channel.connection.client
for _ in eventloop(client, limit=limit, timeout=timeout,
ignore_timeouts=ignore_timeouts):
try:
yield acc.popleft()
except IndexError:
pass
def migrate_tasks(source,
dest,
timeout=1.0,
app=None,
migrate=None,
callback=None):
state = State()
app = app_or_default(app)
def update_state(body, message):
state.count += 1
producer = app.amqp.TaskProducer(dest)
if migrate is None:
migrate = partial(migrate_task, producer)
consumer = app.amqp.TaskConsumer(source)
consumer.register_callback(update_state)
if callback is not None:
callback = partial(callback, state)
consumer.register_callback(callback)
consumer.register_callback(migrate)
# declare all queues on the new broker.
for queue in consumer.queues:
queue(producer.channel).declare()
try:
_, mcount, _ = queue(consumer.channel).queue_declare(passive=True)
if mcount:
state.total_apx += mcount
except source.channel_errors + (StdChannelError, ):
pass
# start migrating messages.
with consumer:
try:
for _ in eventloop(source, timeout=timeout): # pragma: no cover
pass
except socket.timeout:
return
def drain_consumer(consumer,
limit=1,
timeout=None,
callbacks=None,
ignore_timeouts=True):
acc = deque()
def on_message(body, message):
acc.append((body, message))
consumer.callbacks = [on_message] + (callbacks or [])
with consumer:
client = consumer.channel.connection.client
for _ in eventloop(client,
limit=limit,
timeout=timeout,
ignore_timeouts=ignore_timeouts):
try:
yield acc.popleft()
except IndexError:
pass
def drain_events(self, **kwargs):
for _ in eventloop(self.connection, **kwargs):
pass
def drain_events(self, **kwargs):
for _ in eventloop(self.connection, **kwargs):
pass
