RabbitMQQueueSensor

Complexity

Total Complexity

19

Size/Duplication

Total Lines	88
Duplicated Lines	0 %

9 Methods

Rating	Name	Duplication	Size	Complexity
A	cleanup()	0	3	2
A	_deserialize_body()	0	12	3
A	add_trigger()	0	2	1
A	callback()	0	8	1
A	remove_trigger()	0	2	1
B	setup()	0	17	5
A	run()	0	5	1
A	__init__()	0	20	4
A	update_trigger()	0	2	1

import json
import pickle

import eventlet
import pika

from pika.credentials import PlainCredentials

from st2reactor.sensor.base import Sensor

DESERIALIZATION_FUNCTIONS = {
    'json': json.loads,
    'pickle': pickle.loads
}


class RabbitMQQueueSensor(Sensor):
    """Sensor which monitors a RabbitMQ queue for new messages

    This is a RabbitMQ Queue sensor i.e. it works on the simplest RabbitMQ
    messaging model as described in https://www.rabbitmq.com/tutorials/tutorial-one-python.html.

    It is capable of simultaneously consuming from multiple queues. Each message is
    dispatched to stackstorm as a `rabbitmq.new_message` TriggerInstance.
    """
    def __init__(self, sensor_service, config=None):
        super(RabbitMQQueueSensor, self).__init__(sensor_service=sensor_service, config=config)

        self.host = self._config['sensor_config']['host']
        self.username = self._config['sensor_config']['username']
        self.password = self._config['sensor_config']['password']

        queue_sensor_config = self._config['sensor_config']['rabbitmq_queue_sensor']
        self.queues = queue_sensor_config['queues']
        if not isinstance(self.queues, list):
            self.queues = [self.queues]
        self.deserialization_method = queue_sensor_config['deserialization_method']

        supported_methods = DESERIALIZATION_FUNCTIONS.keys()
        if self.deserialization_method and self.deserialization_method not in supported_methods:
            raise ValueError('Invalid deserialization method specified: %s' %
                             (self.deserialization_method))

        self.conn = None
        self.channel = None

    def run(self):
        # run in an eventlet in-order to yield correctly
        gt = eventlet.spawn(self.channel.start_consuming)
        # wait else the sensor will quit
        gt.wait()

    def cleanup(self):
        if self.conn:
            self.conn.close()

    def setup(self):
        if self.username and self.password:
            credentials = PlainCredentials(username=self.username, password=self.password)
            connection_params = pika.ConnectionParameters(host=self.host, credentials=credentials)
        else:
            connection_params = pika.ConnectionParameters(host=self.host)

        self.conn = pika.BlockingConnection(connection_params)
        self.channel = self.conn.channel()
        self.channel.basic_qos(prefetch_count=1)

        # Setup Qs for listening
        for queue in self.queues:
            self.channel.queue_declare(queue=queue, durable=True)
            cb = lambda ch, method, properties, body: self.callback(
                ch, method, properties, body, queue)
            self.channel.basic_consume(cb, queue=queue)

    def callback(self, ch, method, properties, body, queue):
        body = self._deserialize_body(body=body)
        payload = {"queue": queue, "body": body}

        try:
            self._sensor_service.dispatch(trigger="rabbitmq.new_message", payload=payload)
        finally:
            self.channel.basic_ack(delivery_tag=method.delivery_tag)

    def update_trigger(self, trigger):
        pass

    def add_trigger(self, trigger):
        pass

    def remove_trigger(self, trigger):
        pass

    def _deserialize_body(self, body):
        if not self.deserialization_method:
            return body

        deserialization_func = DESERIALIZATION_FUNCTIONS[self.deserialization_method]

        try:
            body = deserialization_func(body)
        except Exception:
            pass

        return body