import pika
import socket
import requests
from pprint import pformat
from cis_interface.drivers.Driver import Driver
from cis_interface.drivers.IODriver import maxMsgSize
from cis_interface.config import cis_cfg
from cis_interface import backwards
[docs]class RMQDriver(Driver):
r"""Class for handling basic RabbitMQ communications.
Args:
name (str): The name of the driver.
queue (str, optional): Name of the queue that messages will be
received from. If and empty string, the queue is exclusive to this
connection. Defaults to ''.
routing_key (str, optional): Routing key that should be used when the
queue is bound. If None, the queue name is used. Defaults to None.
user (str, optional): RabbitMQ server username. Defaults to config
option 'user' in section 'rmq'.
host (str, optional): RabbitMQ server host. Defaults to config option
'host' in section 'rmq' if it exists and the output of
socket.gethostname() if it does not.
vhost (str, optional): RabbitMQ server virtual host. Defaults to
config option 'vhost' in section 'rmq'.
passwd (str, optional): RabbitMQ server password. Defaults to
config option 'password' in section 'rmq'.
exchange (str, optional): RabbitMQ exchange. Defaults to 'namespace'
attribute which is set from the config option 'namespace' in the
section 'rmq'.
exclusive (bool, optional): If True, the queue that is created can
only be used by this driver. Defaults to False. If a queue
name is not provided, it is assumed exclusive.
Attributes:
user (str): RabbitMQ server username.
passwd (str): RabbitMQ server password.
host (str): RabbitMQ server host.
vhost (str): RabbitMQ server virtual host.
exchange (str): RabbitMQ exchange name.
connection (:class:`pika.Connection`): RabbitMQ connection.
channel (:class:`pika.Channel`): RabbitMQ channel.
queue (str): Name of the queue that messages will be received
from. If an empty string, the queue is exclusive to this
connection.
routing_key (str): Routing key that should be used when the queue is
bound. If None, the queue name is used.
times_connected (int): Number of times the connection has been
established/re-established.
"""
def __init__(self, name, queue='', routing_key=None, **kwargs):
kwattr = ['user', 'passwd', 'host', 'vhost', 'exchange', 'exclusive']
kwargs_attr = {k: kwargs.pop(k, None) for k in kwattr}
super(RMQDriver, self).__init__(name, **kwargs)
self.debug()
self.user = cis_cfg.get('rmq', 'user')
self.host = cis_cfg.get('rmq', 'host', socket.gethostname())
self.vhost = cis_cfg.get('rmq', 'vhost')
self.passwd = cis_cfg.get('rmq', 'password')
self.exchange = self.namespace
self.exclusive = False
for k in kwattr:
if kwargs_attr[k] is not None:
setattr(self, k, kwargs_attr.pop(k))
# if getattr(self, k) is None: # pragma: debug
# raise Exception(("%s not provided and corresponding " +
# "environment variable is not set.") % k)
self.connection = None
self.channel = None
self.queue = queue
self.routing_key = routing_key
self.consumer_tag = ""
self._opening = False
self._closing = False
self.times_connected = 0
self.setDaemon(True)
self._q_obj = None
# def __del__(self):
# self.debug('~')
# if self.connection is not None:
# self.connection.close()
# self.connection = None
# except:
# self.debug("::__del__(): exception")
# DRIVER FUNCTIONALITY
[docs] def start(self):
r"""Start the driver. Waiting for connection."""
self._opening = True
super(RMQDriver, self).start()
timeout = self.timeout
while True:
with self.lock:
if not self._opening or timeout <= 0:
break
self.sleep()
timeout -= self.sleeptime
with self.lock:
if self._opening: # pragma: debug
raise RuntimeError("Connection never finished opening.")
[docs] def run(self):
r"""Run the driver. Connect to the connection and begin the IO loop."""
super(RMQDriver, self).run()
self.debug("::run")
self.connect()
self.connection.ioloop.start()
self.debug("::run returns")
[docs] def terminate(self):
r"""Terminate the driver by closing the RabbitMQ connection."""
with self.lock:
if self._closing: # pragma: debug
return # Don't close more than once
self.debug("::terminate")
timeout = self.timeout
while True:
if not self._opening or timeout <= 0:
break
else: # pragma: debug
self.debug('Waiting for connection to open before terminating')
# if self.connection is None:
# self.connection.add_timeout(self.terminate(), self.sleeptime)
# return
# while self.connection is None:
self.sleep()
timeout -= self.sleeptime
self.debug('::terminate: Closing connection')
self.stop_communication()
# Only needed if ioloop is stopped prior to closing the connection?
# (e.g. keyboard interupt)
# if self.connection:
# self.connection.ioloop.start()
super(RMQDriver, self).terminate()
self.debug('::terminate returns')
[docs] def get_message_stats(self):
r"""Return message stats from the server."""
hoststr = self.host
if self.host == '/':
hoststr = '%2f'
url = 'http://%s:%s/api/%s/%s/%s' % (
hoststr, 15672, 'queues', '%2f', self.queue)
res = requests.get(url, auth=(self.user, self.passwd))
jdata = res.json()
if isinstance(jdata, dict):
qdata = jdata.get('message_stats', '')
else:
qdata = ''
return qdata
[docs] def printStatus(self):
r"""Print the driver status."""
self.debug('::printStatus')
super(RMQDriver, self).printStatus()
qdata = self.get_message_stats()
if qdata:
qdata = pformat(qdata)
self.display(": server info:\n%s", qdata)
# RMQ PROPERTIES
[docs] def on_nmsg_request(self, method_frame):
r"""Actions to perform once the queue is declared for message count."""
with self.lock:
self._n_msg = 0
if method_frame:
self._n_msg = method_frame.method.message_count
@property
def n_rmq_msg(self):
r"""int: Number of messages in the queue."""
timeout = self.timeout
self._n_msg = None
with self.lock:
self.channel.queue_declare(self.on_nmsg_request,
queue=self.queue,
exclusive=self.exclusive,
auto_delete=True,
passive=True)
# Wait for queue to be declared passively
while True:
with self.lock:
if self._n_msg is not None or timeout <= 0:
break
self.sleep()
timeout -= self.sleeptime
# Return result
n_msg = 0
if self._n_msg is not None:
n_msg = self._n_msg
return n_msg
@property
def creds(self):
r""":class:`pika.credentials.PlainCredentials`: Server credentials."""
return pika.PlainCredentials(self.user, self.passwd)
@property
def connection_parameters(self):
r""":class:`pika.connection.ConnectionParameters`: Connection
parameters."""
kws = dict(credentials=self.creds,
heartbeat_interval=120,
connection_attempts=3)
if self.host is not None:
kws['host'] = self.host
return pika.ConnectionParameters(**kws)
@property
def is_open(self):
r"""bool: True if connection ready for messages, False otherwise."""
if self.channel is None or self.connection is None:
return False
if self.channel.is_open:
if not self.channel.is_closing:
return True
return False
@property
def is_stable(self):
r"""bool: True if the connection ready for messages and not about to
close. False otherwise."""
if self.is_open and not self._closing:
return True
return False
# CONNECTION
[docs] def connect(self):
r"""Establish the connection."""
self.times_connected += 1
self.connection = pika.SelectConnection(
self.connection_parameters,
on_open_callback=self.on_connection_open,
on_open_error_callback=self.on_connection_open_error,
stop_ioloop_on_close=False)
[docs] def on_connection_open(self, unused_connection):
r"""Actions that must be taken when the connection is opened.
Add the close connection callback and open the RabbitMQ channel."""
self.debug('::Connection opened')
self.connection.add_on_close_callback(self.on_connection_closed)
self.open_channel()
[docs] def on_connection_open_error(self, unused_connection): # pragma: debug
r"""Actions that must be taken when the connection fails to open."""
self.debug('::Connection could not be opened')
self.terminate()
raise Exception('Could not connect.')
[docs] def on_connection_closed(self, connection, reply_code, reply_text):
r"""Actions that must be taken when the connection is closed. Set the
channel to None. If the connection is meant to be closing, stop the
IO loop. Otherwise, wait 5 seconds and try to reconnect."""
with self.lock:
self.debug('::on_connection_closed code %d %s', reply_code,
reply_text)
if self._closing or reply_code == 200:
self.connection.ioloop.stop()
self.connection = None
self._closing = False
else:
self.warn('Connection closed, reopening in %f seconds: (%s) %s',
self.sleeptime, reply_code, reply_text)
self.connection.add_timeout(self.sleeptime, self.reconnect)
[docs] def reconnect(self):
r"""Try to re-establish a connection and resume a new IO loop."""
self.debug('::reconnect()')
# This is the old connection IOLoop instance, stop its ioloop
with self.lock:
self.connection.ioloop.stop()
if not self._closing:
# Create a new connection
self.connect()
# There is now a new connection, needs a new ioloop to run
self.connection.ioloop.start()
# CHANNEL
[docs] def open_channel(self):
r"""Open a RabbitMQ channel."""
self.debug('::Creating a new channel')
self.connection.channel(on_open_callback=self.on_channel_open)
[docs] def on_channel_open(self, channel):
r"""Actions to perform after a channel is opened. Add the channel
close callback and setup the exchange."""
self.debug('::Channel opened')
self.channel = channel
self.channel.add_on_close_callback(self.on_channel_closed)
self.setup_exchange(self.exchange)
[docs] def on_channel_closed(self, channel, reply_code, reply_text):
r"""Actions to perform when the channel is closed. Close the
connection."""
with self.lock:
self.debug('::channel %i was closed: (%s) %s',
channel, reply_code, reply_text)
self.channel = None
self.connection.close()
# EXCHANGE
[docs] def setup_exchange(self, exchange_name):
r"""Setup the exchange."""
self.debug('::Declaring exchange %s', exchange_name)
self.channel.exchange_declare(self.on_exchange_declareok,
exchange=exchange_name, auto_delete=True)
[docs] def on_exchange_declareok(self, unused_frame):
r"""Actions to perform once an exchange is succesfully declared.
Set up the queue."""
self.debug('::Exchange declared')
self.setup_queue(self.queue)
# QUEUE
[docs] def setup_queue(self, queue_name):
r"""Set up the message queue."""
self.debug('::Declaring queue %s', queue_name)
if not queue_name:
self.exclusive = True
self.channel.queue_declare(self.on_queue_declareok,
queue=queue_name,
exclusive=self.exclusive,
auto_delete=True)
[docs] def on_queue_declareok(self, method_frame):
r"""Actions to perform once the queue is succesfully declared. Bind
the queue."""
self.debug('::Binding')
self._q_obj = method_frame.method
self.queue = method_frame.method.queue
if self.routing_key is None:
self.routing_key = self.queue
self.channel.queue_bind(self.on_bindok,
exchange=self.exchange,
queue=self.queue,
routing_key=self.routing_key)
[docs] def on_bindok(self, unused_frame):
r"""Actions to perform once the queue is succesfully bound. Start
consuming messages."""
self.debug('::Queue bound')
with self.lock:
self._opening = False
self.start_communication()
[docs] def purge_queue(self):
r"""Remove all messages from the associated queue."""
with self.lock:
if self._closing or not self.is_open: # pragma: debug
return
if self.channel:
self.channel.queue_purge(queue=self.queue)
# GENERAL
[docs] def start_communication(self, **kwargs):
r"""Start sending/receiving messages."""
pass
[docs] def stop_communication(self, remove_queue=True, **kwargs):
r"""Stop sending/receiving messages."""
self.debug('::stop_communication')
with self.lock:
self._closing = True
if self.is_open:
if remove_queue:
self.debug('::stop_communication: unbinding queue')
self.display('Unbinding queue')
self.channel.queue_unbind(queue=self.queue,
exchange=self.exchange)
self.channel.queue_delete(queue=self.queue)
self.debug('::stop_communication: closing channel')
self.channel.close()
else:
self._closing = False
timeout = self.timeout
while True:
if not self._closing or timeout <= 0:
break
self.debug('::stop_commmunication: waiting for connection to close')
self.sleep()
timeout -= self.sleeptime
# UTILITIES
[docs] def rmq_send(self, data):
r"""Send a message smaller than maxMsgSize to the RMQ queue.
Args:
data (str): The message to be sent.
Returns:
bool: True if the message was sent succesfully. False otherwise.
"""
backwards.assert_bytes(data)
with self.lock:
if self._closing:
return False
self.debug("::send %d", len(data))
assert(len(data) <= maxMsgSize)
if not self.channel: # pragma: debug
self.debug("::send %d NO CHANNEL", len(data))
return False
try:
self.channel.basic_publish(
exchange=self.exchange, routing_key=self.queue,
body=data, mandatory=True)
except Exception as e: # pragma: debug
self.warn("::send %d : exception %s: %s",
len(data), type(e), e)
return False
return True
[docs] def rmq_send_nolimit(self, data):
r"""Send a message smaller than maxMsgSize to the RMQ queue.
Args:
data (str): The message to be sent.
Returns:
bool: True if the message was sent succesfully. False otherwise.
"""
self.debug("::send_nolimit %d", len(data))
prev = 0
ret = self.rmq_send(backwards.unicode2bytes("%ld" % len(data)))
if ret:
while prev < len(data):
next = min(prev + maxMsgSize, len(data))
ret = self.rmq_send(data[prev:next])
prev = next
if not ret: # pragma: debug
break
return ret