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)
