python 使用`pika.BlockingConnection`的RabbitMQ客户端，主要用于发布。

Posted 2021-05-09

import logging
import pika
import json


logger = logging.getLogger(__name__)


class PikaClient(object):
    """A RabbitMQ client using `BlockingConnection`, mainly for publishing.

    It will reconnect and send message again while connection closed by server.
    """
    def __init__(self,
                 credentials=None,
                 host=None,
                 virtual_host=None,
                 exchange=None,
                 routing_key='',
                 heartbeat_interval=60,
                 socket_timeout=2,
                 connection_attempts=3):
        self.exchange = exchange
        self.routing_key = routing_key
        self.connection = None
        self.channel = None
        self._connection_parameters = pika.ConnectionParameters(
            host=host,
            virtual_host=virtual_host,
            credentials=pika.PlainCredentials(*credentials),
            heartbeat_interval=heartbeat_interval,
            socket_timeout=socket_timeout,
            connection_attempts=connection_attempts)
        self._connect()

    def _connect(self):
        self.connection = pika.BlockingConnection(self._connection_parameters)
        self.channel = self.connection.channel()

    def _publish(self, data):
        return self.channel.basic_publish(exchange=self.exchange,
                                          routing_key=self.routing_key,
                                          body=json.dumps(data))

    def publish(self, *args, **kwargs):
        try:
            return self._publish(*args, **kwargs)
        except pika.exceptions.ConnectionClosed:
            logger.info('PikaClient: Connection closed, try reconnecting.')
            self._connect()
            return self._publish(*args, **kwargs)
        except IOError as e:
            logger.info('IOError: {0}'.format(e.message))
            self._connect()
            return self._publish(*args, **kwargs)


if __name__ == '__main__':
    # test code goes below
    logging.basicConfig(level=logging.INFO)
    client = PikaClient(credentials=('L', 'R'),
                        host='x.x.x.x',
                        exchange='test_exchange')

    def test_send_mail(tag):
        client.publish({
            "type": "email",
            "email": ["a@b.c", "[e-mail]"],
            "content-type": "text",
            "subject": "[{0}] Just for test!".format(tag),
            "content": "OOPS...WHAT",
        })