class RabbitMqHandler(object): def __init__(self, connection, exchange_name, type='direct', durable=True): self._logger = logging.getLogger(__name__) try: self._connection = Connection(connection) self._producer = Producer(self._connection) self._task_exchange = Exchange(name=exchange_name, type=type, durable=durable) except Exception: self._logger.info( 'badly formated token %s', auth).exception('Unable to activate the producer') raise def errback(exc, interval): self._logger.info('Error: %r', exc, exc_info=1) self._logger.info('Retry in %s seconds.', interval) def publish(self, payload, routing_key=None, serializer=None): publish = self._connection.ensure(self._producer, self._producer.publish, errback=self.errback, max_retries=3) publish(payload, serializer=serializer, exchange=self._task_exchange, declare=[self._task_exchange], routing_key=routing_key)
class RabbitMqHandler(object): def __init__(self, connection, exchange_name, type='direct', durable=True): self._logger = logging.getLogger(__name__) try: self._connection = Connection(connection) self._producer = Producer(self._connection) self._task_exchange = Exchange(name=exchange_name, type=type, durable=durable) except Exception: self._logger.exception('Impossible to establish the connection') raise def errback(self, exc, interval): self._logger.info('Error: %r', exc, exc_info=1) self._logger.info('Retry in %s seconds.', interval) def publish(self, payload, routing_key=None, serializer=None): try: publish = self._connection.ensure(self._producer, self._producer.publish, errback=self.errback, max_retries=3) publish(payload, serializer=serializer, exchange=self._task_exchange, declare=[self._task_exchange], routing_key=routing_key) except Exception as exc: self._logger.exception('Error occurred when publishing: %r', exc)
def publish_job_action(**message_kwargs): # create a connection to Pulse # TODO: use pulse_conn once we've combined PULSE_URI and PULSE_URL connection = Connection(settings.PULSE_URI) # build up a Producer object from which we can construct the # publish_message function producer = Producer(channel=connection, exchange=Exchange( name="exchange/treeherder/v1/job-actions", type='topic', durable=True, delivery_mode='persistent'), auto_declare=True) publish_message = connection.ensure(producer, producer.publish, max_retries=3) # validate kwargs against the job action schema schemas = load_schemas() schema = "https://treeherder.mozilla.org/schemas/v1/job-action-message.json#" jsonschema.validate(message_kwargs, schemas[schema]) # build the routing key from the message's kwargs routing_key = "{}.{}.{}.".format( message_kwargs["build_system_type"], message_kwargs["project"], message_kwargs["action"], ) publish_message(body=json.dumps(message_kwargs), routing_key=routing_key, content_type='application/json')
def postLogs(logcache): '''post logs to rabbitmq expects a queue object from the multiprocessing library looks for a list of servers in options.mqservers separated by commas creates connections to each, initializes an exchange and a producer and randomly chooses one to publish incoming messages to. ''' mqproducers = list() canQuit = False logger.info('starting message queue posting process') # connect and declare the message queue/kombu objects. # with a list of producers for every potential message queue server. for server in options.mqservers.split(','): connString = 'amqp://{0}:{1}@{2}:{3}//'.format(options.mquser, options.mqpassword, server, options.mqport) mqConn = Connection(connString) eventTaskExchange = Exchange(name=options.taskexchange, type='direct', durable=True) eventTaskExchange(mqConn).declare() mqproducer = mqConn.Producer(serializer='json') ensurePublish = mqConn.ensure(mqproducer, mqproducer.publish, max_retries=10) mqproducers.append(ensurePublish) while True: try: # see if we have anything to post # waiting a bit to not end until we are told we can stop. postdata = logcache.get(True, 1) if postdata is None: # signalled from parent process that it's ok to stop. logcache.task_done() canQuit = True elif len(postdata) > 0: # post to eventtask exchange try: publisher = random.choice(mqproducers) publisher(postdata, exchange=eventTaskExchange, routing_key=options.taskexchange) except Exception as e: logger.error('Exception while posting message: %r' % e) logcache.task_done() except Empty as e: if canQuit: logger.info('shutting down message queue publisher') break logger.info('{0} done'.format('log posting task'))
def postLogs(logcache): '''post logs to rabbitmq expects a queue object from the multiprocessing library looks for a list of servers in options.mqservers separated by commas creates connections to each, initializes an exchange and a producer and randomly chooses one to publish incoming messages to. ''' mqproducers=list() canQuit=False logger.info('starting message queue posting process') #connect and declare the message queue/kombu objects. #with a list of producers for every potential message queue server. for server in options.mqservers.split(','): connString='amqp://{0}:{1}@{2}:{3}//'.format(options.mquser,options.mqpassword,server,options.mqport) mqConn=Connection(connString) eventTaskExchange=Exchange(name=options.taskexchange,type='direct',durable=True) eventTaskExchange(mqConn).declare() mqproducer = mqConn.Producer(serializer='json') ensurePublish=mqConn.ensure(mqproducer,mqproducer.publish,max_retries=10) mqproducers.append(ensurePublish) while True: try: #see if we have anything to post #waiting a bit to not end until we are told we can stop. postdata=logcache.get(False,30) if postdata is None: #signalled from parent process that it's ok to stop. logcache.task_done() canQuit=True elif len(postdata)>0: #post to eventtask exchange try: publisher=random.choice(mqproducers) publisher(postdata,exchange=eventTaskExchange,routing_key=options.taskexchange) except Exception as e: logger.error('Exception while posting message: %r'%e) logcache.task_done() except Empty as e: if canQuit: logger.info('shutting down message queue publisher') break logger.info('{0} done'.format('log posting task'))
def broadcastAttacker(attacker): ''' send this attacker info to our message queue ''' try: connString = 'amqp://{0}:{1}@{2}:{3}/{4}'.format(options.mquser, options.mqpassword, options.mqserver, options.mqport, options.mqvhost) if options.mqprotocol == 'amqps': mqSSL = True else: mqSSL = False mqConn = Connection(connString, ssl=mqSSL) alertExchange = Exchange( name=options.alertexchange, type='topic', durable=True) alertExchange(mqConn).declare() mqproducer = mqConn.Producer(serializer='json') logger.debug('Kombu configured') except Exception as e: logger.error('Exception while configuring kombu for alerts: {0}'.format(e)) try: # generate an 'alert' structure for this attacker: mqAlert = dict(severity='NOTICE', category='attacker') if 'datecreated' in attacker.keys(): mqAlert['utctimestamp'] = attacker['datecreated'].isoformat() mqAlert['summary'] = 'New Attacker: {0} events: {1}, alerts: {2}'.format(attacker['indicators'], attacker['eventscount'], attacker['alertscount']) logger.debug(mqAlert) ensurePublish = mqConn.ensure( mqproducer, mqproducer.publish, max_retries=10) ensurePublish( mqAlert, exchange=alertExchange, routing_key=options.routingkey ) except Exception as e: logger.error('Exception while publishing attacker: {0}'.format(e))
def broadcastAttacker(attacker): ''' send this attacker info to our message queue ''' try: connString = 'amqp://{0}:{1}@{2}:{3}/{4}'.format(options.mquser, options.mqpassword, options.mqserver, options.mqport, options.mqvhost) if options.mqprotocol == 'amqps': mqSSL = True else: mqSSL = False mqConn = Connection(connString, ssl=mqSSL) alertExchange = Exchange( name=options.alertexchange, type='topic', durable=True) alertExchange(mqConn).declare() mqproducer = mqConn.Producer(serializer='json') logger.debug('Kombu configured') except Exception as e: logger.error('Exception while configuring kombu for alerts: {0}'.format(e)) try: # generate an 'alert' structure for this attacker: mqAlert = dict(severity='NOTICE', category='attacker') if 'datecreated' in attacker: mqAlert['utctimestamp'] = attacker['datecreated'].isoformat() mqAlert['summary'] = 'New Attacker: {0} events: {1}, alerts: {2}'.format(attacker['indicators'], attacker['eventscount'], attacker['alertscount']) logger.debug(mqAlert) ensurePublish = mqConn.ensure( mqproducer, mqproducer.publish, max_retries=10) ensurePublish( mqAlert, exchange=alertExchange, routing_key=options.routingkey ) except Exception as e: logger.error('Exception while publishing attacker: {0}'.format(e))
class Propaganda(object): def __init__(self, exchange_name, broker_url=None, key_prefix=''): if not broker_url: if settings.PROPAGANDA_BROKER_URL: broker_url = settings.PROPAGANDA_BROKER_URL else: raise Exception( "You have to define PROPAGANDA_BROKER_URL variable in your django settings" ) self._connection = Connection(broker_url) self._exchange = Exchange(exchange_name, type="topic", channel=self._connection) self._exchange.declare() self._key_prefix = key_prefix self._payload = {} self._subscriptions = {} def __enter__(self): return self def __exit__(self, type, value, traceback): self.connection.release() @property def connection(self): return self._connection @property def exchange(self): return self._exchange @property def key_prefix(self): return self._key_prefix @property def payload(self): return self._payload def subscribe(self, binding_key, reconnect_timeout=10, on_exception=None): if binding_key in self._subscriptions and not self._subscriptions[ binding_key].is_alive(): self._subscriptions[binding_key].start() elif binding_key not in self._subscriptions: self._subscriptions[binding_key] = Subscription( self._connection, self._exchange, binding_key, key_prefix=self._key_prefix, reconnect_timeout=reconnect_timeout, on_exception=on_exception) return self._subscriptions[binding_key] def unsubscribe(self, binding_key): if binding_key in self._subscriptions: self._subscriptions[binding_key].stop() def _publish(self, routing_key, payload, **kwargs): payload.update(self.payload) kwargs['exchange'] = self._exchange kwargs['routing_key'] = '{0}{1}'.format(self._key_prefix, routing_key) with producers[self._connection].acquire(block=True) as producer: publish = self._connection.ensure(producer, producer.publish, max_retries=3) try: publish(payload, **kwargs) except OSError as e: logger.error("Could not publish '{0}': {1}".format( kwargs['routing_key'], e)) else: logger.debug("Published '{0}'".format(kwargs['routing_key'])) def publish(self, routing_key, payload=None, block=True, **kwargs): if payload is None: payload = {} elif not isinstance(payload, dict): logger.error('payload parameter must be a dictionary') raise TypeError("payload parameter must be a dictionary") if block: self._publish(routing_key, payload, **kwargs) else: t = Thread(target=self._publish, args=(routing_key, payload), kwargs=kwargs) t.start()
class test_Connection: def setup(self): self.conn = Connection(port=5672, transport=Transport) def test_establish_connection(self): conn = self.conn assert not conn.connected conn.connect() assert conn.connected assert conn.connection.connected assert conn.host == 'localhost:5672' channel = conn.channel() assert channel.open assert conn.drain_events() == 'event' _connection = conn.connection conn.close() assert not _connection.connected assert isinstance(conn.transport, Transport) def test_reuse_connection(self): conn = self.conn assert conn.connect() is conn.connection is conn.connect() def test_connect_no_transport_options(self): conn = self.conn conn._ensure_connection = Mock() conn.connect() # ensure_connection must be called to return immidiately # and fail with transport exception conn._ensure_connection.assert_called_with( max_retries=1, reraise_as_library_errors=False) def test_connect_transport_options(self): conn = self.conn conn.transport_options = { 'max_retries': 1, 'interval_start': 2, 'interval_step': 3, 'interval_max': 4, 'ignore_this': True } conn._ensure_connection = Mock() conn.connect() # connect() is ignoring transport options # ensure_connection must be called to return immidiately # and fail with transport exception conn._ensure_connection.assert_called_with( max_retries=1, reraise_as_library_errors=False) def test_multiple_urls(self): conn1 = Connection('amqp://foo;amqp://bar') assert conn1.hostname == 'foo' assert conn1.alt == ['amqp://foo', 'amqp://bar'] conn2 = Connection(['amqp://foo', 'amqp://bar']) assert conn2.hostname == 'foo' assert conn2.alt == ['amqp://foo', 'amqp://bar'] def test_collect(self): connection = Connection('memory://') trans = connection._transport = Mock(name='transport') _collect = trans._collect = Mock(name='transport._collect') _close = connection._close = Mock(name='connection._close') connection.declared_entities = Mock(name='decl_entities') uconn = connection._connection = Mock(name='_connection') connection.collect() _close.assert_not_called() _collect.assert_called_with(uconn) connection.declared_entities.clear.assert_called_with() assert trans.client is None assert connection._transport is None assert connection._connection is None def test_prefer_librabbitmq_over_amqp_when_available(self): with patch('kombu.connection.supports_librabbitmq', return_value=True): connection = Connection('amqp://') assert connection.transport_cls == 'librabbitmq' def test_select_amqp_when_librabbitmq_is_not_available(self): with patch('kombu.connection.supports_librabbitmq', return_value=False): connection = Connection('amqp://') assert connection.transport_cls == 'amqp' def test_collect_no_transport(self): connection = Connection('memory://') connection._transport = None connection._do_close_self = Mock() connection._do_close_transport = Mock() connection.collect() connection._do_close_self.assert_called_with() connection._do_close_transport.assert_called_with() connection._do_close_self.side_effect = socket.timeout() connection.collect() def test_collect_transport_gone(self): connection = Connection('memory://') uconn = connection._connection = Mock(name='conn._conn') trans = connection._transport = Mock(name='transport') collect = trans._collect = Mock(name='transport._collect') def se(conn): connection._transport = None collect.side_effect = se connection.collect() collect.assert_called_with(uconn) assert connection._transport is None def test_uri_passthrough(self): transport = Mock(name='transport') with patch('kombu.connection.get_transport_cls') as gtc: gtc.return_value = transport transport.can_parse_url = True with patch('kombu.connection.parse_url') as parse_url: c = Connection('foo+mysql://some_host') assert c.transport_cls == 'foo' parse_url.assert_not_called() assert c.hostname == 'mysql://some_host' assert c.as_uri().startswith('foo+') with patch('kombu.connection.parse_url') as parse_url: c = Connection('mysql://some_host', transport='foo') assert c.transport_cls == 'foo' parse_url.assert_not_called() assert c.hostname == 'mysql://some_host' c = Connection('pyamqp+sqlite://some_host') assert c.as_uri().startswith('pyamqp+') def test_ensure_connection_on_error(self): c = Connection('amqp://A;amqp://B') with patch('kombu.connection.retry_over_time') as rot: c.ensure_connection() rot.assert_called() args = rot.call_args[0] cb = args[4] intervals = iter([1, 2, 3, 4, 5]) assert cb(KeyError(), intervals, 0) == 0 assert cb(KeyError(), intervals, 1) == 1 assert cb(KeyError(), intervals, 2) == 0 assert cb(KeyError(), intervals, 3) == 2 assert cb(KeyError(), intervals, 4) == 0 assert cb(KeyError(), intervals, 5) == 3 assert cb(KeyError(), intervals, 6) == 0 assert cb(KeyError(), intervals, 7) == 4 errback = Mock() c.ensure_connection(errback=errback) args = rot.call_args[0] cb = args[4] assert cb(KeyError(), intervals, 0) == 0 errback.assert_called() def test_supports_heartbeats(self): c = Connection(transport=Mock) c.transport.implements.heartbeats = False assert not c.supports_heartbeats def test_is_evented(self): c = Connection(transport=Mock) c.transport.implements.asynchronous = False assert not c.is_evented def test_register_with_event_loop(self): c = Connection(transport=Mock) loop = Mock(name='loop') c.register_with_event_loop(loop) c.transport.register_with_event_loop.assert_called_with( c.connection, loop, ) def test_manager(self): c = Connection(transport=Mock) assert c.manager is c.transport.manager def test_copy(self): c = Connection('amqp://example.com') assert copy(c).info() == c.info() def test_copy_multiples(self): c = Connection('amqp://A.example.com;amqp://B.example.com') assert c.alt d = copy(c) assert d.alt == c.alt def test_switch(self): c = Connection('amqp://foo') c._closed = True c.switch('redis://example.com//3') assert not c._closed assert c.hostname == 'example.com' assert c.transport_cls == 'redis' assert c.virtual_host == '/3' def test_maybe_switch_next(self): c = Connection('amqp://foo;redis://example.com//3') c.maybe_switch_next() assert not c._closed assert c.hostname == 'example.com' assert c.transport_cls == 'redis' assert c.virtual_host == '/3' def test_maybe_switch_next_no_cycle(self): c = Connection('amqp://foo') c.maybe_switch_next() assert not c._closed assert c.hostname == 'foo' assert c.transport_cls, ('librabbitmq', 'pyamqp' in 'amqp') def test_switch_without_uri_identifier(self): c = Connection('amqp://foo') assert c.hostname == 'foo' assert c.transport_cls, ('librabbitmq', 'pyamqp' in 'amqp') c._closed = True c.switch('example.com') assert not c._closed assert c.hostname == 'example.com' assert c.transport_cls, ('librabbitmq', 'pyamqp' in 'amqp') def test_heartbeat_check(self): c = Connection(transport=Transport) c.transport.heartbeat_check = Mock() c.heartbeat_check(3) c.transport.heartbeat_check.assert_called_with(c.connection, rate=3) def test_completes_cycle_no_cycle(self): c = Connection('amqp://') assert c.completes_cycle(0) assert c.completes_cycle(1) def test_completes_cycle(self): c = Connection('amqp://a;amqp://b;amqp://c') assert not c.completes_cycle(0) assert not c.completes_cycle(1) assert c.completes_cycle(2) def test_get_heartbeat_interval(self): self.conn.transport.get_heartbeat_interval = Mock(name='ghi') assert (self.conn.get_heartbeat_interval() is self.conn.transport.get_heartbeat_interval.return_value) self.conn.transport.get_heartbeat_interval.assert_called_with( self.conn.connection) def test_supports_exchange_type(self): self.conn.transport.implements.exchange_type = {'topic'} assert self.conn.supports_exchange_type('topic') assert not self.conn.supports_exchange_type('fanout') def test_qos_semantics_matches_spec(self): qsms = self.conn.transport.qos_semantics_matches_spec = Mock() assert self.conn.qos_semantics_matches_spec is qsms.return_value qsms.assert_called_with(self.conn.connection) def test__enter____exit__(self): conn = self.conn context = conn.__enter__() assert context is conn conn.connect() assert conn.connection.connected conn.__exit__() assert conn.connection is None conn.close() # again def test_close_survives_connerror(self): class _CustomError(Exception): pass class MyTransport(Transport): connection_errors = (_CustomError, ) def close_connection(self, connection): raise _CustomError('foo') conn = Connection(transport=MyTransport) conn.connect() conn.close() assert conn._closed def test_close_when_default_channel(self): conn = self.conn conn._default_channel = Mock() conn._close() conn._default_channel.close.assert_called_with() def test_auto_reconnect_default_channel(self): # tests GH issue: #1208 # Tests that default_channel automatically reconnects when connection # closed c = Connection('memory://') c._closed = True with patch.object(c, '_connection_factory', side_effect=c._connection_factory) as cf_mock: c.default_channel cf_mock.assert_called_once_with() def test_close_when_default_channel_close_raises(self): class Conn(Connection): @property def connection_errors(self): return (KeyError, ) conn = Conn('memory://') conn._default_channel = Mock() conn._default_channel.close.side_effect = KeyError() conn._close() conn._default_channel.close.assert_called_with() def test_revive_when_default_channel(self): conn = self.conn defchan = conn._default_channel = Mock() conn.revive(Mock()) defchan.close.assert_called_with() assert conn._default_channel is None def test_ensure_connection(self): assert self.conn.ensure_connection() def test_ensure_success(self): def publish(): return 'foobar' ensured = self.conn.ensure(None, publish) assert ensured() == 'foobar' def test_ensure_failure(self): class _CustomError(Exception): pass def publish(): raise _CustomError('bar') ensured = self.conn.ensure(None, publish) with pytest.raises(_CustomError): ensured() def test_ensure_connection_failure(self): class _ConnectionError(Exception): pass def publish(): raise _ConnectionError('failed connection') self.conn.transport.connection_errors = (_ConnectionError, ) ensured = self.conn.ensure(self.conn, publish) with pytest.raises(OperationalError): ensured() def test_autoretry(self): myfun = Mock() self.conn.transport.connection_errors = (KeyError, ) def on_call(*args, **kwargs): myfun.side_effect = None raise KeyError('foo') myfun.side_effect = on_call insured = self.conn.autoretry(myfun) insured() myfun.assert_called() def test_SimpleQueue(self): conn = self.conn q = conn.SimpleQueue('foo') assert q.channel is conn.default_channel chan = conn.channel() q2 = conn.SimpleQueue('foo', channel=chan) assert q2.channel is chan def test_SimpleBuffer(self): conn = self.conn q = conn.SimpleBuffer('foo') assert q.channel is conn.default_channel chan = conn.channel() q2 = conn.SimpleBuffer('foo', channel=chan) assert q2.channel is chan def test_SimpleQueue_with_parameters(self): conn = self.conn q = conn.SimpleQueue('foo', True, {'durable': True}, {'x-queue-mode': 'lazy'}, { 'durable': True, 'type': 'fanout', 'delivery_mode': 'persistent' }) assert q.queue.exchange.type == 'fanout' assert q.queue.exchange.durable assert not q.queue.exchange.auto_delete delivery_mode_code = q.queue.exchange.PERSISTENT_DELIVERY_MODE assert q.queue.exchange.delivery_mode == delivery_mode_code assert q.queue.queue_arguments['x-queue-mode'] == 'lazy' assert q.queue.durable assert not q.queue.auto_delete def test_SimpleBuffer_with_parameters(self): conn = self.conn q = conn.SimpleBuffer('foo', True, {'durable': True}, {'x-queue-mode': 'lazy'}, { 'durable': True, 'type': 'fanout', 'delivery_mode': 'persistent' }) assert q.queue.exchange.type == 'fanout' assert q.queue.exchange.durable assert q.queue.exchange.auto_delete delivery_mode_code = q.queue.exchange.PERSISTENT_DELIVERY_MODE assert q.queue.exchange.delivery_mode == delivery_mode_code assert q.queue.queue_arguments['x-queue-mode'] == 'lazy' assert q.queue.durable assert q.queue.auto_delete def test_Producer(self): conn = self.conn assert isinstance(conn.Producer(), Producer) assert isinstance(conn.Producer(conn.default_channel), Producer) def test_Consumer(self): conn = self.conn assert isinstance(conn.Consumer(queues=[]), Consumer) assert isinstance( conn.Consumer(queues=[], channel=conn.default_channel), Consumer) def test__repr__(self): assert repr(self.conn) def test__reduce__(self): x = pickle.loads(pickle.dumps(self.conn)) assert x.info() == self.conn.info() def test_channel_errors(self): class MyTransport(Transport): channel_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) assert conn.channel_errors == (KeyError, ValueError) def test_connection_errors(self): class MyTransport(Transport): connection_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) assert conn.connection_errors == (KeyError, ValueError) def test_multiple_urls_hostname(self): conn = Connection(['example.com;amqp://example.com']) assert conn.as_uri() == 'amqp://*****:*****@example.com:5672//' conn = Connection(['example.com', 'amqp://example.com']) assert conn.as_uri() == 'amqp://*****:*****@example.com:5672//' conn = Connection('example.com;example.com;') assert conn.as_uri() == 'amqp://*****:*****@example.com:5672//'
#rabbit_url = "amqp://*****:*****@134.221.121.63:5672/vhost1" #no proxy conn = Connection(rabbit_url, transport_options={'confirm_publish': True}) producer = Producer(conn) exchange = Exchange("ex2", type="topic", durable=True) queues = Queue("hello7", exchange, routing_key="hello7", durable=True) def errback(exc, interval): logging.error('Error: %r', exc, exc_info=1) logging.info('Retry in %s seconds.', interval) publish = conn.ensure(producer, producer.publish, errback=errback, max_retries=3) msg_body = 'START' print str(datetime.now()) + " publishing " + msg_body publish(msg_body, routing_key='hello7', exchange=queues.exchange, declare=[queues]) conn.release() t = 1 while t < 101: publish = conn.ensure(producer, producer.publish, errback=errback, max_retries=3)
class KombuMessenger: """ Sends messages via Kombu. """ def __init__(self, queueHost, queueName, id, hostname, pid, type): """ Initializer. """ self._queueHost = queueHost self._queueName = queueName self._id = id self._hostname = hostname self._pid = pid self._type = type self._connection = Connection('pyamqp://*****:*****@%s:5672//' % self._queueHost) self._connection.ensure_connection() self._exchange = Exchange(self._queueName, type='direct') self._queue = Queue(self._queueName, self._exchange, routing_key=self._queueName) self._producer = self._connection.Producer() self._publish = self._connection.ensure(self._producer, self._producer.publish, max_retries=3) # end def def __del__(self): """ Finalizer. """ self._connection.close() # end def def __str__(self): """ Gets the string representation of this object. @return: the string representation of this object. @rtype: str """ return 'connection: "%s", id: "%s", queueName: "%s", hostname: "%s", pid: "%s", type: "%s"' % ( self._connection, self._id, self._queueName, self._hostname, self._pid, self._type) # end def def send(self, chunk): """ Send stream chunk with JSON descriptor. """ context = { 'id': self._id, 'datetime': datetime.isoformat(datetime.now()), 'hostname': self._hostname, 'pid': self._pid, 'type': self._type, 'chunk': chunk } #contextStr = json.dumps(context) self._publish(context, routing_key=self._queueName, declare=[self._queue])
# Declare the Task Exchange for events # delivery_mode=1 is fast/auto-ack messages, 2 is require ack. # mozdef default exchange is: eventtask, routing key is also: eventtask eventTaskExchange = Exchange(name='eventtask', type='direct', durable=True, delivery_mode=1) eventTaskExchange(mqConn).declare() mqproducer = mqConn.Producer(serializer='json') # make an event event = dict() # best practice is to send an ISO formatted timestamp # so upstream can tell the source time zone event['timestamp'] = pytz.timezone('UTC').localize( datetime.utcnow()).isoformat() event['summary'] = 'just a test, only a test' event['category'] = 'testing' event['severity'] = 'INFO' event['processid'] = os.getpid() event['processname'] = sys.argv[0] event['tags'] = list() event['tags'].append('test') event['details'] = dict() event['details']['sourceipaddress'] = '1.2.3.4' # publish it to rabbit mq ensurePublish = mqConn.ensure(mqproducer, mqproducer.publish, max_retries=10) ensurePublish(event, exchange=eventTaskExchange, routing_key='eventtask')
class test_Connection: def setup(self): self.conn = Connection(port=5672, transport=Transport) def test_establish_connection(self): conn = self.conn conn.connect() assert conn.connection.connected assert conn.host == 'localhost:5672' channel = conn.channel() assert channel.open assert conn.drain_events() == 'event' _connection = conn.connection conn.close() assert not _connection.connected assert isinstance(conn.transport, Transport) def test_multiple_urls(self): conn1 = Connection('amqp://foo;amqp://bar') assert conn1.hostname == 'foo' assert conn1.alt == ['amqp://foo', 'amqp://bar'] conn2 = Connection(['amqp://foo', 'amqp://bar']) assert conn2.hostname == 'foo' assert conn2.alt == ['amqp://foo', 'amqp://bar'] def test_collect(self): connection = Connection('memory://') trans = connection._transport = Mock(name='transport') _collect = trans._collect = Mock(name='transport._collect') _close = connection._close = Mock(name='connection._close') connection.declared_entities = Mock(name='decl_entities') uconn = connection._connection = Mock(name='_connection') connection.collect() _close.assert_not_called() _collect.assert_called_with(uconn) connection.declared_entities.clear.assert_called_with() assert trans.client is None assert connection._transport is None assert connection._connection is None def test_collect_no_transport(self): connection = Connection('memory://') connection._transport = None connection._do_close_self = Mock() connection._do_close_transport = Mock() connection.collect() connection._do_close_self.assert_called_with() connection._do_close_transport.assert_called_with() connection._do_close_self.side_effect = socket.timeout() connection.collect() def test_collect_transport_gone(self): connection = Connection('memory://') uconn = connection._connection = Mock(name='conn._conn') trans = connection._transport = Mock(name='transport') collect = trans._collect = Mock(name='transport._collect') def se(conn): connection._transport = None collect.side_effect = se connection.collect() collect.assert_called_with(uconn) assert connection._transport is None def test_uri_passthrough(self): transport = Mock(name='transport') with patch('kombu.connection.get_transport_cls') as gtc: gtc.return_value = transport transport.can_parse_url = True with patch('kombu.connection.parse_url') as parse_url: c = Connection('foo+mysql://some_host') assert c.transport_cls == 'foo' parse_url.assert_not_called() assert c.hostname == 'mysql://some_host' assert c.as_uri().startswith('foo+') with patch('kombu.connection.parse_url') as parse_url: c = Connection('mysql://some_host', transport='foo') assert c.transport_cls == 'foo' parse_url.assert_not_called() assert c.hostname == 'mysql://some_host' c = Connection('pyamqp+sqlite://some_host') assert c.as_uri().startswith('pyamqp+') def test_default_ensure_callback(self): with patch('kombu.connection.logger') as logger: c = Connection(transport=Mock) c._default_ensure_callback(KeyError(), 3) logger.error.assert_called() def test_ensure_connection_on_error(self): c = Connection('amqp://A;amqp://B') with patch('kombu.connection.retry_over_time') as rot: c.ensure_connection() rot.assert_called() args = rot.call_args[0] cb = args[4] intervals = iter([1, 2, 3, 4, 5]) assert cb(KeyError(), intervals, 0) == 0 assert cb(KeyError(), intervals, 1) == 1 assert cb(KeyError(), intervals, 2) == 0 assert cb(KeyError(), intervals, 3) == 2 assert cb(KeyError(), intervals, 4) == 0 assert cb(KeyError(), intervals, 5) == 3 assert cb(KeyError(), intervals, 6) == 0 assert cb(KeyError(), intervals, 7) == 4 errback = Mock() c.ensure_connection(errback=errback) args = rot.call_args[0] cb = args[4] assert cb(KeyError(), intervals, 0) == 0 errback.assert_called() def test_supports_heartbeats(self): c = Connection(transport=Mock) c.transport.implements.heartbeats = False assert not c.supports_heartbeats def test_is_evented(self): c = Connection(transport=Mock) c.transport.implements.async = False assert not c.is_evented def test_register_with_event_loop(self): c = Connection(transport=Mock) loop = Mock(name='loop') c.register_with_event_loop(loop) c.transport.register_with_event_loop.assert_called_with( c.connection, loop, ) def test_manager(self): c = Connection(transport=Mock) assert c.manager is c.transport.manager def test_copy(self): c = Connection('amqp://example.com') assert copy(c).info() == c.info() def test_copy_multiples(self): c = Connection('amqp://A.example.com;amqp://B.example.com') assert c.alt d = copy(c) assert d.alt == c.alt def test_switch(self): c = Connection('amqp://foo') c._closed = True c.switch('redis://example.com//3') assert not c._closed assert c.hostname == 'example.com' assert c.transport_cls == 'redis' assert c.virtual_host == '/3' def test_maybe_switch_next(self): c = Connection('amqp://foo;redis://example.com//3') c.maybe_switch_next() assert not c._closed assert c.hostname == 'example.com' assert c.transport_cls == 'redis' assert c.virtual_host == '/3' def test_maybe_switch_next_no_cycle(self): c = Connection('amqp://foo') c.maybe_switch_next() assert not c._closed assert c.hostname == 'foo' assert c.transport_cls, ('librabbitmq', 'pyamqp' in 'amqp') def test_heartbeat_check(self): c = Connection(transport=Transport) c.transport.heartbeat_check = Mock() c.heartbeat_check(3) c.transport.heartbeat_check.assert_called_with(c.connection, rate=3) def test_completes_cycle_no_cycle(self): c = Connection('amqp://') assert c.completes_cycle(0) assert c.completes_cycle(1) def test_completes_cycle(self): c = Connection('amqp://a;amqp://b;amqp://c') assert not c.completes_cycle(0) assert not c.completes_cycle(1) assert c.completes_cycle(2) def test_get_heartbeat_interval(self): self.conn.transport.get_heartbeat_interval = Mock(name='ghi') assert (self.conn.get_heartbeat_interval() is self.conn.transport.get_heartbeat_interval.return_value) self.conn.transport.get_heartbeat_interval.assert_called_with( self.conn.connection) def test_supports_exchange_type(self): self.conn.transport.implements.exchange_type = {'topic'} assert self.conn.supports_exchange_type('topic') assert not self.conn.supports_exchange_type('fanout') def test_qos_semantics_matches_spec(self): qsms = self.conn.transport.qos_semantics_matches_spec = Mock() assert self.conn.qos_semantics_matches_spec is qsms.return_value qsms.assert_called_with(self.conn.connection) def test__enter____exit__(self): conn = self.conn context = conn.__enter__() assert context is conn conn.connect() assert conn.connection.connected conn.__exit__() assert conn.connection is None conn.close() # again def test_close_survives_connerror(self): class _CustomError(Exception): pass class MyTransport(Transport): connection_errors = (_CustomError,) def close_connection(self, connection): raise _CustomError('foo') conn = Connection(transport=MyTransport) conn.connect() conn.close() assert conn._closed def test_close_when_default_channel(self): conn = self.conn conn._default_channel = Mock() conn._close() conn._default_channel.close.assert_called_with() def test_close_when_default_channel_close_raises(self): class Conn(Connection): @property def connection_errors(self): return (KeyError,) conn = Conn('memory://') conn._default_channel = Mock() conn._default_channel.close.side_effect = KeyError() conn._close() conn._default_channel.close.assert_called_with() def test_revive_when_default_channel(self): conn = self.conn defchan = conn._default_channel = Mock() conn.revive(Mock()) defchan.close.assert_called_with() assert conn._default_channel is None def test_ensure_connection(self): assert self.conn.ensure_connection() def test_ensure_success(self): def publish(): return 'foobar' ensured = self.conn.ensure(None, publish) assert ensured() == 'foobar' def test_ensure_failure(self): class _CustomError(Exception): pass def publish(): raise _CustomError('bar') ensured = self.conn.ensure(None, publish) with pytest.raises(_CustomError): ensured() def test_ensure_connection_failure(self): class _ConnectionError(Exception): pass def publish(): raise _ConnectionError('failed connection') self.conn.transport.connection_errors = (_ConnectionError,) ensured = self.conn.ensure(self.conn, publish) with pytest.raises(OperationalError): ensured() def test_autoretry(self): myfun = Mock() self.conn.transport.connection_errors = (KeyError,) def on_call(*args, **kwargs): myfun.side_effect = None raise KeyError('foo') myfun.side_effect = on_call insured = self.conn.autoretry(myfun) insured() myfun.assert_called() def test_SimpleQueue(self): conn = self.conn q = conn.SimpleQueue('foo') assert q.channel is conn.default_channel chan = conn.channel() q2 = conn.SimpleQueue('foo', channel=chan) assert q2.channel is chan def test_SimpleBuffer(self): conn = self.conn q = conn.SimpleBuffer('foo') assert q.channel is conn.default_channel chan = conn.channel() q2 = conn.SimpleBuffer('foo', channel=chan) assert q2.channel is chan def test_Producer(self): conn = self.conn assert isinstance(conn.Producer(), Producer) assert isinstance(conn.Producer(conn.default_channel), Producer) def test_Consumer(self): conn = self.conn assert isinstance(conn.Consumer(queues=[]), Consumer) assert isinstance( conn.Consumer(queues=[], channel=conn.default_channel), Consumer) def test__repr__(self): assert repr(self.conn) def test__reduce__(self): x = pickle.loads(pickle.dumps(self.conn)) assert x.info() == self.conn.info() def test_channel_errors(self): class MyTransport(Transport): channel_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) assert conn.channel_errors == (KeyError, ValueError) def test_connection_errors(self): class MyTransport(Transport): connection_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) assert conn.connection_errors == (KeyError, ValueError)
class test_Connection(Case): def setup(self): self.conn = Connection(port=5672, transport=Transport) def test_establish_connection(self): conn = self.conn conn.connect() self.assertTrue(conn.connection.connected) self.assertEqual(conn.host, 'localhost:5672') channel = conn.channel() self.assertTrue(channel.open) self.assertEqual(conn.drain_events(), 'event') _connection = conn.connection conn.close() self.assertFalse(_connection.connected) self.assertIsInstance(conn.transport, Transport) def test_multiple_urls(self): conn1 = Connection('amqp://foo;amqp://bar') self.assertEqual(conn1.hostname, 'foo') self.assertListEqual(conn1.alt, ['amqp://foo', 'amqp://bar']) conn2 = Connection(['amqp://foo', 'amqp://bar']) self.assertEqual(conn2.hostname, 'foo') self.assertListEqual(conn2.alt, ['amqp://foo', 'amqp://bar']) def test_collect(self): connection = Connection('memory://') trans = connection._transport = Mock(name='transport') _collect = trans._collect = Mock(name='transport._collect') _close = connection._close = Mock(name='connection._close') connection.declared_entities = Mock(name='decl_entities') uconn = connection._connection = Mock(name='_connection') connection.collect() self.assertFalse(_close.called) _collect.assert_called_with(uconn) connection.declared_entities.clear.assert_called_with() self.assertIsNone(trans.client) self.assertIsNone(connection._transport) self.assertIsNone(connection._connection) def test_collect_no_transport(self): connection = Connection('memory://') connection._transport = None connection._do_close_self = Mock() connection._do_close_transport = Mock() connection.collect() connection._do_close_self.assert_called_with() connection._do_close_transport.assert_called_with() connection._do_close_self.side_effect = socket.timeout() connection.collect() def test_collect_transport_gone(self): connection = Connection('memory://') uconn = connection._connection = Mock(name='conn._conn') trans = connection._transport = Mock(name='transport') collect = trans._collect = Mock(name='transport._collect') def se(conn): connection._transport = None collect.side_effect = se connection.collect() collect.assert_called_with(uconn) self.assertIsNone(connection._transport) def test_uri_passthrough(self): transport = Mock(name='transport') with patch('kombu.connection.get_transport_cls') as gtc: gtc.return_value = transport transport.can_parse_url = True with patch('kombu.connection.parse_url') as parse_url: c = Connection('foo+mysql://some_host') self.assertEqual(c.transport_cls, 'foo') self.assertFalse(parse_url.called) self.assertEqual(c.hostname, 'mysql://some_host') self.assertTrue(c.as_uri().startswith('foo+')) with patch('kombu.connection.parse_url') as parse_url: c = Connection('mysql://some_host', transport='foo') self.assertEqual(c.transport_cls, 'foo') self.assertFalse(parse_url.called) self.assertEqual(c.hostname, 'mysql://some_host') c = Connection('pyamqp+sqlite://some_host') self.assertTrue(c.as_uri().startswith('pyamqp+')) def test_default_ensure_callback(self): with patch('kombu.connection.logger') as logger: c = Connection(transport=Mock) c._default_ensure_callback(KeyError(), 3) self.assertTrue(logger.error.called) def test_ensure_connection_on_error(self): c = Connection('amqp://A;amqp://B') with patch('kombu.connection.retry_over_time') as rot: c.ensure_connection() self.assertTrue(rot.called) args = rot.call_args[0] cb = args[4] intervals = iter([1, 2, 3, 4, 5]) self.assertEqual(cb(KeyError(), intervals, 0), 0) self.assertEqual(cb(KeyError(), intervals, 1), 1) self.assertEqual(cb(KeyError(), intervals, 2), 0) self.assertEqual(cb(KeyError(), intervals, 3), 2) self.assertEqual(cb(KeyError(), intervals, 4), 0) self.assertEqual(cb(KeyError(), intervals, 5), 3) self.assertEqual(cb(KeyError(), intervals, 6), 0) self.assertEqual(cb(KeyError(), intervals, 7), 4) errback = Mock() c.ensure_connection(errback=errback) args = rot.call_args[0] cb = args[4] self.assertEqual(cb(KeyError(), intervals, 0), 0) self.assertTrue(errback.called) def test_supports_heartbeats(self): c = Connection(transport=Mock) c.transport.implements.heartbeats = False self.assertFalse(c.supports_heartbeats) def test_is_evented(self): c = Connection(transport=Mock) c.transport.implements. async = False self.assertFalse(c.is_evented) def test_register_with_event_loop(self): c = Connection(transport=Mock) loop = Mock(name='loop') c.register_with_event_loop(loop) c.transport.register_with_event_loop.assert_called_with( c.connection, loop, ) def test_manager(self): c = Connection(transport=Mock) self.assertIs(c.manager, c.transport.manager) def test_copy(self): c = Connection('amqp://example.com') self.assertEqual(copy(c).info(), c.info()) def test_copy_multiples(self): c = Connection('amqp://A.example.com;amqp://B.example.com') self.assertTrue(c.alt) d = copy(c) self.assertEqual(d.alt, c.alt) def test_switch(self): c = Connection('amqp://foo') c._closed = True c.switch('redis://example.com//3') self.assertFalse(c._closed) self.assertEqual(c.hostname, 'example.com') self.assertEqual(c.transport_cls, 'redis') self.assertEqual(c.virtual_host, '/3') def test_maybe_switch_next(self): c = Connection('amqp://foo;redis://example.com//3') c.maybe_switch_next() self.assertFalse(c._closed) self.assertEqual(c.hostname, 'example.com') self.assertEqual(c.transport_cls, 'redis') self.assertEqual(c.virtual_host, '/3') def test_maybe_switch_next_no_cycle(self): c = Connection('amqp://foo') c.maybe_switch_next() self.assertFalse(c._closed) self.assertEqual(c.hostname, 'foo') self.assertIn(c.transport_cls, ('librabbitmq', 'pyamqp', 'amqp')) def test_heartbeat_check(self): c = Connection(transport=Transport) c.transport.heartbeat_check = Mock() c.heartbeat_check(3) c.transport.heartbeat_check.assert_called_with(c.connection, rate=3) def test_completes_cycle_no_cycle(self): c = Connection('amqp://') self.assertTrue(c.completes_cycle(0)) self.assertTrue(c.completes_cycle(1)) def test_completes_cycle(self): c = Connection('amqp://a;amqp://b;amqp://c') self.assertFalse(c.completes_cycle(0)) self.assertFalse(c.completes_cycle(1)) self.assertTrue(c.completes_cycle(2)) def test__enter____exit__(self): conn = self.conn context = conn.__enter__() self.assertIs(context, conn) conn.connect() self.assertTrue(conn.connection.connected) conn.__exit__() self.assertIsNone(conn.connection) conn.close() # again def test_close_survives_connerror(self): class _CustomError(Exception): pass class MyTransport(Transport): connection_errors = (_CustomError, ) def close_connection(self, connection): raise _CustomError('foo') conn = Connection(transport=MyTransport) conn.connect() conn.close() self.assertTrue(conn._closed) def test_close_when_default_channel(self): conn = self.conn conn._default_channel = Mock() conn._close() conn._default_channel.close.assert_called_with() def test_close_when_default_channel_close_raises(self): class Conn(Connection): @property def connection_errors(self): return (KeyError, ) conn = Conn('memory://') conn._default_channel = Mock() conn._default_channel.close.side_effect = KeyError() conn._close() conn._default_channel.close.assert_called_with() def test_revive_when_default_channel(self): conn = self.conn defchan = conn._default_channel = Mock() conn.revive(Mock()) defchan.close.assert_called_with() self.assertIsNone(conn._default_channel) def test_ensure_connection(self): self.assertTrue(self.conn.ensure_connection()) def test_ensure_success(self): def publish(): return 'foobar' ensured = self.conn.ensure(None, publish) self.assertEqual(ensured(), 'foobar') def test_ensure_failure(self): class _CustomError(Exception): pass def publish(): raise _CustomError('bar') ensured = self.conn.ensure(None, publish) with self.assertRaises(_CustomError): ensured() def test_ensure_connection_failure(self): class _ConnectionError(Exception): pass def publish(): raise _ConnectionError('failed connection') self.conn.transport.connection_errors = (_ConnectionError, ) ensured = self.conn.ensure(self.conn, publish) with self.assertRaises(_ConnectionError): ensured() def test_autoretry(self): myfun = Mock() self.conn.transport.connection_errors = (KeyError, ) def on_call(*args, **kwargs): myfun.side_effect = None raise KeyError('foo') myfun.side_effect = on_call insured = self.conn.autoretry(myfun) insured() self.assertTrue(myfun.called) def test_SimpleQueue(self): conn = self.conn q = conn.SimpleQueue('foo') self.assertIs(q.channel, conn.default_channel) chan = conn.channel() q2 = conn.SimpleQueue('foo', channel=chan) self.assertIs(q2.channel, chan) def test_SimpleBuffer(self): conn = self.conn q = conn.SimpleBuffer('foo') self.assertIs(q.channel, conn.default_channel) chan = conn.channel() q2 = conn.SimpleBuffer('foo', channel=chan) self.assertIs(q2.channel, chan) def test_Producer(self): conn = self.conn self.assertIsInstance(conn.Producer(), Producer) self.assertIsInstance(conn.Producer(conn.default_channel), Producer) def test_Consumer(self): conn = self.conn self.assertIsInstance(conn.Consumer(queues=[]), Consumer) self.assertIsInstance( conn.Consumer(queues=[], channel=conn.default_channel), Consumer) def test__repr__(self): self.assertTrue(repr(self.conn)) def test__reduce__(self): x = pickle.loads(pickle.dumps(self.conn)) self.assertDictEqual(x.info(), self.conn.info()) def test_channel_errors(self): class MyTransport(Transport): channel_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) self.assertTupleEqual(conn.channel_errors, (KeyError, ValueError)) def test_connection_errors(self): class MyTransport(Transport): connection_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) self.assertTupleEqual(conn.connection_errors, (KeyError, ValueError))
#ssl or not mqConn = Connection(connString, ssl=False) # Declare the Task Exchange for events # delivery_mode=1 is fast/auto-ack messages, 2 is require ack. # mozdef default exchange is: eventtask, routing key is also: eventtask eventTaskExchange = Exchange(name='eventtask', type='direct', durable=True, delivery_mode=1) eventTaskExchange(mqConn).declare() mqproducer = mqConn.Producer(serializer='json') # make an event event = dict() # best practice is to send an ISO formatted timestamp # so upstream can tell the source time zone event['timestamp'] = pytz.timezone('UTC').localize(datetime.utcnow()).isoformat() event['summary'] = 'just a test, only a test' event['category'] = 'testing' event['severity'] = 'INFO' event['processid']=os.getpid() event['processname']=sys.argv[0] event['tags'] = list() event['tags'].append('test') event['details'] = dict() event['details']['sourceipaddress'] = '1.2.3.4' # publish it to rabbit mq ensurePublish=mqConn.ensure(mqproducer,mqproducer.publish,max_retries=10) ensurePublish(event,exchange=eventTaskExchange,routing_key='eventtask')
# Kombu Connection to message broker kombuConnection = Connection('amqp://*****:*****@rabbit:5672//') # Initialization of Producer producer = Producer(kombuConnection, auto_declare=True) # Kombu error callback def errback(exc, interval): logging.error('Error: %r', exc, exc_info=1) logging.error('Retries in second', interval) # Ensuring broker Connection publish = kombuConnection.ensure(producer, producer.publish, errback=errback, max_retries=3) @celery.task(acks_late=True) def long_task(name): logs_queue = Queue(long_task.request.id, routing_key=long_task.request.id) # Sending Message back to broker publish({'Message': { 'value': 'Executing Long Task' }}, routing_key=logs_queue.routing_key, declare=[logs_queue]) cmd = "sh jobs/longTaskjob.sh {}".format(name) process = Popen(shlex.split(cmd),
class test_Connection(TestCase): def setUp(self): self.conn = Connection(port=5672, transport=Transport) def test_establish_connection(self): conn = self.conn conn.connect() self.assertTrue(conn.connection.connected) self.assertEqual(conn.host, 'localhost:5672') channel = conn.channel() self.assertTrue(channel.open) self.assertEqual(conn.drain_events(), 'event') _connection = conn.connection conn.close() self.assertFalse(_connection.connected) self.assertIsInstance(conn.transport, Transport) def test__enter____exit__(self): conn = self.conn context = conn.__enter__() self.assertIs(context, conn) conn.connect() self.assertTrue(conn.connection.connected) conn.__exit__() self.assertIsNone(conn.connection) conn.close() # again def test_close_survives_connerror(self): class _CustomError(Exception): pass class MyTransport(Transport): connection_errors = (_CustomError, ) def close_connection(self, connection): raise _CustomError('foo') conn = Connection(transport=MyTransport) conn.connect() conn.close() self.assertTrue(conn._closed) def test_close_when_default_channel(self): conn = self.conn conn._default_channel = Mock() conn._close() conn._default_channel.close.assert_called_with() def test_close_when_default_channel_close_raises(self): class Conn(Connection): @property def connection_errors(self): return (KeyError, ) conn = Conn('memory://') conn._default_channel = Mock() conn._default_channel.close.side_effect = KeyError() conn._close() conn._default_channel.close.assert_called_with() def test_revive_when_default_channel(self): conn = self.conn defchan = conn._default_channel = Mock() conn.revive(Mock()) defchan.close.assert_called_with() self.assertIsNone(conn._default_channel) def test_ensure_connection(self): self.assertTrue(self.conn.ensure_connection()) def test_ensure_success(self): def publish(): return 'foobar' ensured = self.conn.ensure(None, publish) self.assertEqual(ensured(), 'foobar') def test_ensure_failure(self): class _CustomError(Exception): pass def publish(): raise _CustomError('bar') ensured = self.conn.ensure(None, publish) with self.assertRaises(_CustomError): ensured() def test_ensure_connection_failure(self): class _ConnectionError(Exception): pass def publish(): raise _ConnectionError('failed connection') self.conn.transport.connection_errors = (_ConnectionError, ) ensured = self.conn.ensure(self.conn, publish) with self.assertRaises(_ConnectionError): ensured() def test_autoretry(self): myfun = Mock() myfun.__name__ = 'test_autoretry' self.conn.transport.connection_errors = (KeyError, ) def on_call(*args, **kwargs): myfun.side_effect = None raise KeyError('foo') myfun.side_effect = on_call insured = self.conn.autoretry(myfun) insured() self.assertTrue(myfun.called) def test_SimpleQueue(self): conn = self.conn q = conn.SimpleQueue('foo') self.assertIs(q.channel, conn.default_channel) chan = conn.channel() q2 = conn.SimpleQueue('foo', channel=chan) self.assertIs(q2.channel, chan) def test_SimpleBuffer(self): conn = self.conn q = conn.SimpleBuffer('foo') self.assertIs(q.channel, conn.default_channel) chan = conn.channel() q2 = conn.SimpleBuffer('foo', channel=chan) self.assertIs(q2.channel, chan) def test_Producer(self): conn = self.conn self.assertIsInstance(conn.Producer(), Producer) self.assertIsInstance(conn.Producer(conn.default_channel), Producer) def test_Consumer(self): conn = self.conn self.assertIsInstance(conn.Consumer(queues=[]), Consumer) self.assertIsInstance( conn.Consumer(queues=[], channel=conn.default_channel), Consumer) def test__repr__(self): self.assertTrue(repr(self.conn)) def test__reduce__(self): x = pickle.loads(pickle.dumps(self.conn)) self.assertDictEqual(x.info(), self.conn.info()) def test_channel_errors(self): class MyTransport(Transport): channel_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) self.assertTupleEqual(conn.channel_errors, (KeyError, ValueError)) def test_connection_errors(self): class MyTransport(Transport): connection_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) self.assertTupleEqual(conn.connection_errors, (KeyError, ValueError))
class Emitter(object): """Helper class for writing robust AMQP producers. An instance of this class will connect to the given AMQP broker, and allow you to send messages to an exchange on that broker with a given routing key, optionally also declaring a queue and and binding it to the exchange with that routing key. Upon connection errors (either upon initial connection or message send) it will retry several times before giving up, and if a connection is interrupted it will recover gracefully on the next attempt to send a message. Short-lived connections ----------------------- To create a short-lived connection which automatically releases/disconnects when you've finished, just use an instance of this class as a context manager; for example:: url = broker_url('guest', 'guest', 'localhost') with Emitter(url, 'ex_name', 'key') as emitter: emitter.send_message({'status': 'success'}, serializer='json') Note the lack of any explicit `connect()` call here: `send_message()` calls it if it hasn't been called already, so you never need to call it explicitly (though you can). Note also the lack of any queue_name parameter in the Emitter constructor call: the message will be sent to the given exchange with the given routing key, but no queue declaration/binding takes place here (compare with the next example). Long-lived connections ---------------------- Alternatively, you can just create an instance of this class and have it live for as long as you need it, in which case the connection will remain open until the instance falls out of scope or you call its `release()` method. To have an instance live for the lifetime of an application, a helpful pattern is to put it in a module-level global and lazily declare it on first use, e.g.:: # Module-level global for the instance; initially nothing here. _emitter = None def send_message(msg, serializer): global _emitter if _emitter is None: url = broker_url('guest', 'guest', 'localhost') _emitter = Emitter(url, 'ex_name', 'key', 'q_name') _emitter.send_message(msg, serializer) Then you can just call this module's `send_message()` function whenever you need to send a message: it will connect to the broker on the first call, keep the connection open between calls, and attempt to reconnect on the next call if the connection ever goes away. Note also that in this example we do pass a queue_name to the Emitter constructor: upon connection, the Emitter will ensure that the named queue is declared, and bound to the given exchange with the given routing key. Customisation via subclassing ----------------------------- You may wish to use a custom subclass to override certain behaviour around error handling. In particular, you can override: * Various `CONN_*` and `SEND_*` class variables, which control connection/send retry behaviour. E.g. `SEND_MAX_RETRIES` defines how many times to retry sending a message before giving up. * The `conn_errback()` callback method, which is called upon an error while connecting. * The `send_errback()` callback method, which is called upon an error while sending a message. * The `errback()` callback method, which (by default) is just what `conn_errback()` and `send_errback()` call - so override just this to change them both. The default version just logs the error. For example, here's a class which retries connecting forever; waits up to 10 seconds between retries; only retries message sends once; and logs the full traceback for connect errors:: class CustomEmitter(Emitter): CONN_MAX_RETRIES = None CONN_INTERVAL_MAX = 10 SEND_MAX_RETRIES = 1 def conn_errback(self, ex, interval): logger.error('Error: %r', ex, exc_info=1) logger.info('Retry in %s seconds.', interval) Note the following edge cases for `*_MAX_RETRIES`: * `None` - for both `CONN_` and `SEND_` this means "try forever". * 0 - for `CONN_` this means "try forever" but for `SEND_` this means "retry once", i.e. the same as `SEND_MAX_RETRIES = 1`; this seems to be a kombu bug. * Negative values - for `CONN_` this means "just try once/no retries", but for `SEND_` it just breaks things (it will never recover the connection); this seems to be a kombu bug. """ CONN_MAX_RETRIES = 2 # Max retries when attempting to connect. CONN_INTERVAL_START = 2 # Seconds to wait between connect retries, initially. CONN_INTERVAL_STEP = 2 # Increase time between connect retries by this amount. CONN_INTERVAL_MAX = 4 # Maximum time between connect retries. SEND_MAX_RETRIES = 3 # Max retries when attempting to send. SEND_INTERVAL_START = 1 # Seconds to wait between send retries, initially. SEND_INTERVAL_STEP = 1 # Increase time between send retries by this amount. SEND_INTERVAL_MAX = 1 # Maximum time between send retries. def __init__(self, url, exchange_name, routing_key, queue_name=None, exchange_type='direct'): logger.debug('Initialising {}'.format(self.__class__.__name__)) self.url = url self.exchange_name = exchange_name self.routing_key = routing_key self.queue_name = queue_name self.exchange_type = exchange_type self._connection = None self._producer = None def conn_errback(self, ex, interval): """Callback called upon connection error.""" self.errback(ex, interval) def send_errback(self, ex, interval): """Callback called upon send error.""" self.errback(ex, interval) def errback(self, ex, interval): """Default callback called upon connection or send error.""" logger.info('Error: {} - {}'.format(ex.__class__.__name__, str(ex))) logger.info('Retry in %s seconds.', interval) def connect(self): """Connect to broker and possibly ensure exchange/queue/routing declared. In case of errors, this will retry connecting several times (controlled by the `CONN_*` class variables), and gracefully recover if possible. If it never succeeds, the exception raised by the final attempt is re-raised. """ logger.debug('Connecting to broker at: {}'.format(self.url)) self._connection = Connection(self.url) # Kombu interprets interval_max incorrectly; work around that. interval_max = self.CONN_INTERVAL_MAX - self.CONN_INTERVAL_STEP self._connection.ensure_connection( errback=self.conn_errback, max_retries=self.CONN_MAX_RETRIES, interval_start=self.CONN_INTERVAL_START, interval_step=self.CONN_INTERVAL_STEP, interval_max=interval_max) exchange = Exchange(name=self.exchange_name, type=self.exchange_type) channel = self._connection.channel() self._producer = Producer(channel=channel, exchange=exchange, routing_key=self.routing_key) if self.queue_name: # Bind/declare queue. queue = Queue(name=self.queue_name, exchange=exchange, routing_key=self.routing_key) queue = queue(channel) # Bind queue logger.debug('Declaring queue {}, on exchange {} at {}'.format( self.queue_name, self.exchange_name, self.url)) queue.declare() def release(self): """Disconnect/release.""" self._producer.release() self._connection.release() def __enter__(self): """Context manager entry: connect to the broker.""" self.connect() return self def __exit__(self, exc_type, exc_value, traceback): """Context manager exit: disconnect/release.""" self.release() def send_message(self, message, serializer='json', headers=None): """Send a message with retries (and connect to broker if necessary). In case of errors, this will retry sending several times (controlled by the `SEND_*` class variables), and gracefully recover if possible. If it never succeeds, the exception raised by the final attempt is re-raised. You can use the headers argument to pass in any custom headers. It is a dictionary {"a-custom-header": "a-custom-value"}. """ logger.debug("Sending message...") if self._producer is None: self.connect() # Kombu interprets interval_max incorrectly; work around that. interval_max = self.SEND_INTERVAL_MAX - self.SEND_INTERVAL_STEP publish = self._connection.ensure( self._producer, self._producer.publish, errback=self.send_errback, max_retries=self.SEND_MAX_RETRIES, interval_start=self.SEND_INTERVAL_START, interval_step=self.SEND_INTERVAL_STEP, interval_max=interval_max) publish(message, serializer=serializer, headers=headers)
except: return '' try: ## mongodb connection if os.getenv('MONGODB_USER'): client = MongoClient('mongodb://{}:{}/'.format(os.getenv('MONGODB_HOST'),os.getenv('MONGODB_PORT')), username=os.getenv('MONGODB_USER'), password=os.getenv('MONGODB_PASS'), authSource='admin', document_class=RawBSONDocument) else: client = MongoClient('mongodb://{}:{}/'.format(os.getenv('MONGODB_HOST'),os.getenv('MONGODB_PORT')), document_class=RawBSONDocument) db = client[os.getenv('MONGODB_DB', 'v2')] collection = db[KEY] ## rabbitmq connection connection = Connection('amqp://{}:{}@{}:{}//'.format(os.getenv('RABBITMQ_DEFAULT_USER'),os.getenv('RABBITMQ_DEFAULT_PASS'),os.getenv('RABBITMQ_HOST'),os.getenv('RABBITMQ_PORT'))) exchange = Exchange(KEY, type='topic') producer = Producer(connection, exchange=exchange) publisher = connection.ensure(producer, producer.publish, max_retries=5) ## redis connection r = redis.StrictRedis(host=os.getenv("REDIS_HOST"), db=int(os.getenv("REDIS_DB")), password=os.getenv("REDIS_PASS")) except: time.sleep(1) sys.exit(1) def createQueues(qKey): return [Queue(qKey, [binding(exchange, routing_key='{}.{}'.format(qKey, action)) for action in ['create','update','delete']])] def create${KEY}(body=None): # noqa: E501 """Create a new ${KEY} # noqa: E501
class AmqpPublisher(object): def __init__(self, _conf): amqp_conf = conf.get(ConfigKeys.AMQP) queue_host = amqp_conf.get(ConfigKeys.HOST) if queue_host is None or len(queue_host.strip()) == 0: return queue_port = amqp_conf.get(ConfigKeys.PORT) queue_vhost = amqp_conf.get(ConfigKeys.VHOST) queue_user = amqp_conf.get(ConfigKeys.USER) queue_pass = amqp_conf.get(ConfigKeys.PASSWORD) queue_host = ';'.join( ['amqp://%s' % host for host in queue_host.split(';')]) queue_exchange = '%s_%s' % (amqp_conf.get( ConfigKeys.EXCHANGE), amqp_conf.get(ConfigKeys.ENVIRONMENT)) queue_name = amqp_conf.get(ConfigKeys.QUEUE) self.exchange = Exchange(queue_exchange, type='direct') self.queue_connection = Connection(hostname=queue_host, port=queue_port, virtual_host=queue_vhost, userid=queue_user, password=queue_pass) self.queue = Queue(queue_name, self.exchange) logger.info('setting up pubsub for host(s) "{}"'.format(queue_host)) def error_callback(self, exc, interval) -> None: logger.warning('could not connect to MQ (interval: %s): %s' % (str(interval), str(exc))) def try_publish(self, message): logger.info('sending with "{}"'.format(str(self.queue_connection))) with producers[self.queue_connection].acquire(block=False) as producer: amqp_publish = self.queue_connection.ensure( producer, producer.publish, errback=self.error_callback, max_retries=3) amqp_publish(message, exchange=self.exchange, declare=[self.exchange, self.queue]) def publish(self, message: dict) -> None: n_tries = 3 current_try = 0 failed = False for current_try in range(n_tries): try: self.try_publish(message) failed = False break except Exception as pe: failed = True logger.error('[%s/%s tries] failed to publish external: %s' % (str(current_try + 1), str(n_tries), str(pe))) logger.exception(traceback.format_exc()) time.sleep(0.1) if failed: raise PublishException() elif current_try > 0: logger.info('published event on attempt {}/{}'.format( str(current_try + 1), str(n_tries))) else: logger.debug('published event with verb {} id {}'.format( message['verb'], message['id']))
class Publisher(KombuConfReader): def __init__(self, config): self._log = log.getLogger() KombuConfReader.__init__(self, config) self.connection = Connection(self.broker_url) try: self._init_amqp() except Exception as exc: self._log.error('Publisher fail in init connection: %s' % exc) raise def _init_amqp(self): """Init AMQP objects after connection""" self.producer = self.connection.Producer() self.exchange = Exchange( self.exchange_name, channel=self.connection.channel(), type=self.exchange_type, durable=self.exchange_is_durable) self.queue = Queue( self.queue_name, self.exchange, channel=self.connection.channel(), durable=self.queue_is_durable, routing_key=self.routing_key, queue_arguments=self.queue_args) # We declare object to broker. this way only we can # ensure to publish to an existing queue and routing_key # AMQP work this way, not a library principle self.exchange.declare() self.queue.declare() def switch_connection(self): """Switch AMQP connection from url to backup_url and vice versa""" self._log.warn('Switching AMQP connection from %s' % self.connection.as_uri()) if (self.connection.hostname in self.broker_url and self.broker_backup_url): self.connection = Connection(self.broker_backup_url) elif self.connection.hostname in self.broker_backup_url: self.connection = Connection(self.broker_url) else: raise URLError('Invalid current URI to switch connection : %s' % self.connection.as_uri()) self._init_amqp() def _publish(self, msg): """Publish message ensuring connection is available""" publish = self.connection.ensure( self.producer, self.producer.publish, max_retries=3) publish(msg, exchange=self.exchange, routing_key=self.routing_key, serializer=self.serializer, compression=self.compression) return True def publish(self, msg): """ must return True/False if message is well publish or not """ try: return self._publish(msg) except Exception as exc: try: self.switch_connection() return self._publish(msg) except Exception as exc: self._log.error('Publish fail when switching connection: %s' % exc) return False
class test_Connection(TestCase): def setUp(self): self.conn = Connection(port=5672, transport=Transport) def test_establish_connection(self): conn = self.conn conn.connect() self.assertTrue(conn.connection.connected) self.assertEqual(conn.host, 'localhost:5672') channel = conn.channel() self.assertTrue(channel.open) self.assertEqual(conn.drain_events(), 'event') _connection = conn.connection conn.close() self.assertFalse(_connection.connected) self.assertIsInstance(conn.transport, Transport) def test__enter____exit__(self): conn = self.conn context = conn.__enter__() self.assertIs(context, conn) conn.connect() self.assertTrue(conn.connection.connected) conn.__exit__() self.assertIsNone(conn.connection) conn.close() # again def test_close_survives_connerror(self): class _CustomError(Exception): pass class MyTransport(Transport): connection_errors = (_CustomError, ) def close_connection(self, connection): raise _CustomError('foo') conn = Connection(transport=MyTransport) conn.connect() conn.close() self.assertTrue(conn._closed) def test_close_when_default_channel(self): conn = self.conn conn._default_channel = Mock() conn._close() conn._default_channel.close.assert_called_with() def test_close_when_default_channel_close_raises(self): class Conn(Connection): @property def connection_errors(self): return (KeyError, ) conn = Conn('memory://') conn._default_channel = Mock() conn._default_channel.close.side_effect = KeyError() conn._close() conn._default_channel.close.assert_called_with() def test_revive_when_default_channel(self): conn = self.conn defchan = conn._default_channel = Mock() conn.revive(Mock()) defchan.close.assert_called_with() self.assertIsNone(conn._default_channel) def test_ensure_connection(self): self.assertTrue(self.conn.ensure_connection()) def test_ensure_success(self): def publish(): return 'foobar' ensured = self.conn.ensure(None, publish) self.assertEqual(ensured(), 'foobar') def test_ensure_failure(self): class _CustomError(Exception): pass def publish(): raise _CustomError('bar') ensured = self.conn.ensure(None, publish) with self.assertRaises(_CustomError): ensured() def test_ensure_connection_failure(self): class _ConnectionError(Exception): pass def publish(): raise _ConnectionError('failed connection') self.conn.transport.connection_errors = (_ConnectionError,) ensured = self.conn.ensure(self.conn, publish) with self.assertRaises(_ConnectionError): ensured() def test_autoretry(self): myfun = Mock() myfun.__name__ = 'test_autoretry' self.conn.transport.connection_errors = (KeyError, ) def on_call(*args, **kwargs): myfun.side_effect = None raise KeyError('foo') myfun.side_effect = on_call insured = self.conn.autoretry(myfun) insured() self.assertTrue(myfun.called) def test_SimpleQueue(self): conn = self.conn q = conn.SimpleQueue('foo') self.assertIs(q.channel, conn.default_channel) chan = conn.channel() q2 = conn.SimpleQueue('foo', channel=chan) self.assertIs(q2.channel, chan) def test_SimpleBuffer(self): conn = self.conn q = conn.SimpleBuffer('foo') self.assertIs(q.channel, conn.default_channel) chan = conn.channel() q2 = conn.SimpleBuffer('foo', channel=chan) self.assertIs(q2.channel, chan) def test_Producer(self): conn = self.conn self.assertIsInstance(conn.Producer(), Producer) self.assertIsInstance(conn.Producer(conn.default_channel), Producer) def test_Consumer(self): conn = self.conn self.assertIsInstance(conn.Consumer(queues=[]), Consumer) self.assertIsInstance(conn.Consumer(queues=[], channel=conn.default_channel), Consumer) def test__repr__(self): self.assertTrue(repr(self.conn)) def test__reduce__(self): x = pickle.loads(pickle.dumps(self.conn)) self.assertDictEqual(x.info(), self.conn.info()) def test_channel_errors(self): class MyTransport(Transport): channel_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) self.assertTupleEqual(conn.channel_errors, (KeyError, ValueError)) def test_connection_errors(self): class MyTransport(Transport): connection_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) self.assertTupleEqual(conn.connection_errors, (KeyError, ValueError))
class Queue(object): """ Manage connections to rabbitmq. """ def __init__(self): self.rabbitmq_host = os.getenv('RABBITMQ_HOST', 'localhost') self.rabbitmq_port = os.getenv('RABBITMQ_PORT', 5672) self.rabbitmq_username = os.getenv('RABBITMQ_USERNAME', 'guest') self.rabbitmq_password = os.getenv('RABBITMQ_PASSWORD', 'guest') self.rabbitmq_exchange = os.getenv('RABBITMQ_EXCHANGE', 'votes') self.queue = gevent.queue.Queue() self.create_connection() def start(self): """ Start the queue service running. """ greenlets = [] try: logging.debug("starting publish greenlet") greenlets.append(gevent.spawn(self.publish_loop)) except: logging.exception("unexpected greenlet exit") gevent.killall(greenlets) raise def create_connection(self): """ Create a connection the the broker. """ self.connection = Connection('amqp://{}:{}@{}:{}//'.format( self.rabbitmq_username, self.rabbitmq_password, self.rabbitmq_host, self.rabbitmq_port )) self.producer = Producer(self.connection) def queue_message(self, message): """ Add a message to the gevent queue to get published. """ logging.info('Adding message to GEVENT queue to be published.') self.queue.put(message) def publish_message(self, message): """ Publish a single message to rabbitmq. """ logging.debug("throwing message on the queue.") publish = self.connection.ensure(self.producer, self.producer.publish, max_retries=3) publish(message, routing_key='votes') def publish_loop(self): """ Listen to the Gevent queue and publish messages. """ logging.warning('starting queue, waiting for events.') for message in self.queue: logging.warning("calling publish_message") self.publish_message(message)
class Publisher(KombuConfReader): def __init__(self, config): self._log = log.getLogger() KombuConfReader.__init__(self, config) self.connection = Connection(self.broker_url) try: self._init_amqp() except Exception as exc: self._log.error('Publisher fail in init connection: %s' % exc) raise def _init_amqp(self): """Init AMQP objects after connection""" self.producer = self.connection.Producer() self.exchange = Exchange(self.exchange_name, channel=self.connection.channel(), type=self.exchange_type, durable=self.exchange_is_durable) self.queue = Queue(self.queue_name, self.exchange, channel=self.connection.channel(), durable=self.queue_is_durable, routing_key=self.routing_key, queue_arguments=self.queue_args) # We declare object to broker. this way only we can # ensure to publish to an existing queue and routing_key # AMQP work this way, not a library principle self.exchange.declare() self.queue.declare() def switch_connection(self): """Switch AMQP connection from url to backup_url and vice versa""" self._log.warn('Switching AMQP connection from %s' % self.connection.as_uri()) if (self.connection.hostname in self.broker_url and self.broker_backup_url): self.connection = Connection(self.broker_backup_url) elif self.connection.hostname in self.broker_backup_url: self.connection = Connection(self.broker_url) else: raise URLError('Invalid current URI to switch connection : %s' % self.connection.as_uri()) self._init_amqp() def _publish(self, msg): """Publish message ensuring connection is available""" publish = self.connection.ensure(self.producer, self.producer.publish, max_retries=3) publish(msg, exchange=self.exchange, routing_key=self.routing_key, serializer=self.serializer, compression=self.compression) return True def publish(self, msg): """ must return True/False if message is well publish or not """ try: return self._publish(msg) except Exception as exc: try: self.switch_connection() return self._publish(msg) except Exception as exc: self._log.error('Publish fail when switching connection: %s' % exc) return False
class test_Connection(Case): def setUp(self): self.conn = Connection(port=5672, transport=Transport) def test_establish_connection(self): conn = self.conn conn.connect() self.assertTrue(conn.connection.connected) self.assertEqual(conn.host, 'localhost:5672') channel = conn.channel() self.assertTrue(channel.open) self.assertEqual(conn.drain_events(), 'event') _connection = conn.connection conn.close() self.assertFalse(_connection.connected) self.assertIsInstance(conn.transport, Transport) def test_multiple_urls(self): conn1 = Connection('amqp://foo;amqp://bar') self.assertEqual(conn1.hostname, 'foo') self.assertListEqual(conn1.alt, ['amqp://foo', 'amqp://bar']) conn2 = Connection(['amqp://foo', 'amqp://bar']) self.assertEqual(conn2.hostname, 'foo') self.assertListEqual(conn2.alt, ['amqp://foo', 'amqp://bar']) def test_collect(self): connection = Connection('memory://') trans = connection._transport = Mock(name='transport') _collect = trans._collect = Mock(name='transport._collect') _close = connection._close = Mock(name='connection._close') connection.declared_entities = Mock(name='decl_entities') uconn = connection._connection = Mock(name='_connection') connection.collect() self.assertFalse(_close.called) _collect.assert_called_with(uconn) connection.declared_entities.clear.assert_called_with() self.assertIsNone(trans.client) self.assertIsNone(connection._transport) self.assertIsNone(connection._connection) def test_collect_no_transport(self): connection = Connection('memory://') connection._transport = None connection._close = Mock() connection.collect() connection._close.assert_called_with() connection._close.side_effect = socket.timeout() connection.collect() def test_collect_transport_gone(self): connection = Connection('memory://') uconn = connection._connection = Mock(name='conn._conn') trans = connection._transport = Mock(name='transport') collect = trans._collect = Mock(name='transport._collect') def se(conn): connection._transport = None collect.side_effect = se connection.collect() collect.assert_called_with(uconn) self.assertIsNone(connection._transport) def test_uri_passthrough(self): transport = Mock(name='transport') with patch('kombu.connection.get_transport_cls') as gtc: gtc.return_value = transport transport.can_parse_url = True with patch('kombu.connection.parse_url') as parse_url: c = Connection('foo+mysql://some_host') self.assertEqual(c.transport_cls, 'foo') self.assertFalse(parse_url.called) self.assertEqual(c.hostname, 'mysql://some_host') self.assertTrue(c.as_uri().startswith('foo+')) with patch('kombu.connection.parse_url') as parse_url: c = Connection('mysql://some_host', transport='foo') self.assertEqual(c.transport_cls, 'foo') self.assertFalse(parse_url.called) self.assertEqual(c.hostname, 'mysql://some_host') c = Connection('pyamqp+sqlite://some_host') self.assertTrue(c.as_uri().startswith('pyamqp+')) def test_default_ensure_callback(self): with patch('kombu.connection.logger') as logger: c = Connection(transport=Mock) c._default_ensure_callback(KeyError(), 3) self.assertTrue(logger.error.called) def test_ensure_connection_on_error(self): c = Connection('amqp://A;amqp://B') with patch('kombu.connection.retry_over_time') as rot: c.ensure_connection() self.assertTrue(rot.called) args = rot.call_args[0] cb = args[4] intervals = iter([1, 2, 3, 4, 5]) self.assertEqual(cb(KeyError(), intervals, 0), 0) self.assertEqual(cb(KeyError(), intervals, 1), 1) self.assertEqual(cb(KeyError(), intervals, 2), 0) self.assertEqual(cb(KeyError(), intervals, 3), 2) self.assertEqual(cb(KeyError(), intervals, 4), 0) self.assertEqual(cb(KeyError(), intervals, 5), 3) self.assertEqual(cb(KeyError(), intervals, 6), 0) self.assertEqual(cb(KeyError(), intervals, 7), 4) errback = Mock() c.ensure_connection(errback=errback) args = rot.call_args[0] cb = args[4] self.assertEqual(cb(KeyError(), intervals, 0), 0) self.assertTrue(errback.called) def test_supports_heartbeats(self): c = Connection(transport=Mock) c.transport.supports_heartbeats = False self.assertFalse(c.supports_heartbeats) def test_is_evented(self): c = Connection(transport=Mock) c.transport.supports_ev = False self.assertFalse(c.is_evented) def test_register_with_event_loop(self): c = Connection(transport=Mock) loop = Mock(name='loop') c.register_with_event_loop(loop) c.transport.register_with_event_loop.assert_called_with( c.connection, loop, ) def test_manager(self): c = Connection(transport=Mock) self.assertIs(c.manager, c.transport.manager) def test_copy(self): c = Connection('amqp://example.com') self.assertEqual(copy(c).info(), c.info()) def test_copy_multiples(self): c = Connection('amqp://A.example.com;amqp://B.example.com') self.assertTrue(c.alt) d = copy(c) self.assertEqual(d.alt, c.alt) def test_switch(self): c = Connection('amqp://foo') c._closed = True c.switch('redis://example.com//3') self.assertFalse(c._closed) self.assertEqual(c.hostname, 'example.com') self.assertEqual(c.transport_cls, 'redis') self.assertEqual(c.virtual_host, '/3') def test_maybe_switch_next(self): c = Connection('amqp://foo;redis://example.com//3') c.maybe_switch_next() self.assertFalse(c._closed) self.assertEqual(c.hostname, 'example.com') self.assertEqual(c.transport_cls, 'redis') self.assertEqual(c.virtual_host, '/3') def test_maybe_switch_next_no_cycle(self): c = Connection('amqp://foo') c.maybe_switch_next() self.assertFalse(c._closed) self.assertEqual(c.hostname, 'foo') self.assertIn(c.transport_cls, ('librabbitmq', 'pyamqp', 'amqp')) def test_heartbeat_check(self): c = Connection(transport=Transport) c.transport.heartbeat_check = Mock() c.heartbeat_check(3) c.transport.heartbeat_check.assert_called_with(c.connection, rate=3) def test_completes_cycle_no_cycle(self): c = Connection('amqp://') self.assertTrue(c.completes_cycle(0)) self.assertTrue(c.completes_cycle(1)) def test_completes_cycle(self): c = Connection('amqp://a;amqp://b;amqp://c') self.assertFalse(c.completes_cycle(0)) self.assertFalse(c.completes_cycle(1)) self.assertTrue(c.completes_cycle(2)) def test__enter____exit__(self): conn = self.conn context = conn.__enter__() self.assertIs(context, conn) conn.connect() self.assertTrue(conn.connection.connected) conn.__exit__() self.assertIsNone(conn.connection) conn.close() # again def test_close_survives_connerror(self): class _CustomError(Exception): pass class MyTransport(Transport): connection_errors = (_CustomError, ) def close_connection(self, connection): raise _CustomError('foo') conn = Connection(transport=MyTransport) conn.connect() conn.close() self.assertTrue(conn._closed) def test_close_when_default_channel(self): conn = self.conn conn._default_channel = Mock() conn._close() conn._default_channel.close.assert_called_with() def test_close_when_default_channel_close_raises(self): class Conn(Connection): @property def connection_errors(self): return (KeyError, ) conn = Conn('memory://') conn._default_channel = Mock() conn._default_channel.close.side_effect = KeyError() conn._close() conn._default_channel.close.assert_called_with() def test_revive_when_default_channel(self): conn = self.conn defchan = conn._default_channel = Mock() conn.revive(Mock()) defchan.close.assert_called_with() self.assertIsNone(conn._default_channel) def test_ensure_connection(self): self.assertTrue(self.conn.ensure_connection()) def test_ensure_success(self): def publish(): return 'foobar' ensured = self.conn.ensure(None, publish) self.assertEqual(ensured(), 'foobar') def test_ensure_failure(self): class _CustomError(Exception): pass def publish(): raise _CustomError('bar') ensured = self.conn.ensure(None, publish) with self.assertRaises(_CustomError): ensured() def test_ensure_connection_failure(self): class _ConnectionError(Exception): pass def publish(): raise _ConnectionError('failed connection') self.conn.transport.connection_errors = (_ConnectionError,) ensured = self.conn.ensure(self.conn, publish) with self.assertRaises(_ConnectionError): ensured() def test_autoretry(self): myfun = Mock() self.conn.transport.connection_errors = (KeyError, ) def on_call(*args, **kwargs): myfun.side_effect = None raise KeyError('foo') myfun.side_effect = on_call insured = self.conn.autoretry(myfun) insured() self.assertTrue(myfun.called) def test_SimpleQueue(self): conn = self.conn q = conn.SimpleQueue('foo') self.assertIs(q.channel, conn.default_channel) chan = conn.channel() q2 = conn.SimpleQueue('foo', channel=chan) self.assertIs(q2.channel, chan) def test_SimpleBuffer(self): conn = self.conn q = conn.SimpleBuffer('foo') self.assertIs(q.channel, conn.default_channel) chan = conn.channel() q2 = conn.SimpleBuffer('foo', channel=chan) self.assertIs(q2.channel, chan) def test_Producer(self): conn = self.conn self.assertIsInstance(conn.Producer(), Producer) self.assertIsInstance(conn.Producer(conn.default_channel), Producer) def test_Consumer(self): conn = self.conn self.assertIsInstance(conn.Consumer(queues=[]), Consumer) self.assertIsInstance(conn.Consumer(queues=[], channel=conn.default_channel), Consumer) def test__repr__(self): self.assertTrue(repr(self.conn)) def test__reduce__(self): x = pickle.loads(pickle.dumps(self.conn)) self.assertDictEqual(x.info(), self.conn.info()) def test_channel_errors(self): class MyTransport(Transport): channel_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) self.assertTupleEqual(conn.channel_errors, (KeyError, ValueError)) def test_connection_errors(self): class MyTransport(Transport): connection_errors = (KeyError, ValueError) conn = Connection(transport=MyTransport) self.assertTupleEqual(conn.connection_errors, (KeyError, ValueError))
class Queue(object): """ Manage connections to rabbitmq. """ def __init__(self): self.rabbitmq_host = os.getenv('RABBITMQ_HOST', 'localhost') self.rabbitmq_port = os.getenv('RABBITMQ_PORT', 5672) self.rabbitmq_username = os.getenv('RABBITMQ_USERNAME', 'guest') self.rabbitmq_password = os.getenv('RABBITMQ_PASSWORD', 'guest') self.rabbitmq_exchange = os.getenv('RABBITMQ_EXCHANGE', 'votes') self.queue = gevent.queue.Queue() self.create_connection() def start(self): """ Start the queue service running. """ greenlets = [] try: logging.debug("starting publish greenlet") greenlets.append(gevent.spawn(self.publish_loop)) except: logging.exception("unexpected greenlet exit") gevent.killall(greenlets) raise def create_connection(self): """ Create a connection the the broker. """ self.connection = Connection('amqp://{}:{}@{}:{}//'.format( self.rabbitmq_username, self.rabbitmq_password, self.rabbitmq_host, self.rabbitmq_port)) self.producer = Producer(self.connection) def queue_message(self, message): """ Add a message to the gevent queue to get published. """ logging.info('Adding message to GEVENT queue to be published.') self.queue.put(message) def publish_message(self, message): """ Publish a single message to rabbitmq. """ logging.debug("throwing message on the queue.") publish = self.connection.ensure(self.producer, self.producer.publish, max_retries=3) publish(message, routing_key='votes') def publish_loop(self): """ Listen to the Gevent queue and publish messages. """ logging.warning('starting queue, waiting for events.') for message in self.queue: logging.warning("calling publish_message") self.publish_message(message)