class RabbitMQPublisher(object):
def __init__(self, source_id, network_id, model_key=None, run_key=None):
broker_url = get_broker_url(model_key)
self.conn = Connection(broker_url)
self.conn.connect()
self.channel = 'oa-{source_id}-{network_id}-{model_key}'.format(
source_id=source_id,
network_id=network_id,
model_key=model_key or environ.get(constants.MODEL_KEY)
)
if run_key is not None:
self.channel += '-{}'.format(run_key)
self.producer = self.conn.Producer(serializer='json')
# set up the Exchange, Queue, and Producer
media_exchange = Exchange('media', 'direct', durable=True)
video_queue = Queue('video', exchange=media_exchange, routing_key='video')
# publish updates
def publish(self, payload):
self.producer.publish(self.channel, json.dumps(payload))
def close(self):
self.conn.release()
def start_celery_task(task: str, arguments: List[object],
task_id: int) -> None:
"""
Generic function to start a task through celery
Arguments:
task {string} -- Name of the task to start.
arguments {List[object]} -- List of arguments to be supplied.
id {int} -- Integer ID to uniquely identify the string.
"""
task_exchange = Exchange("", type="direct")
connection = Connection(RABBIT_MQ_ADDRESS)
producer = Producer(
exchange=task_exchange,
channel=connection.channel(),
routing_key="celery",
serializer="json",
)
# is the task ID necessary here?
payload = {
"id": task_id,
"task": task,
"args": arguments,
"kwargs": {},
"retries": 0,
}
producer.publish(
json.dumps(payload, default=json_util.default),
content_type="application/json",
content_encoding="utf-8",
)
connection.release()
def check_server_access(sconf):
servers = ["broker server", "results server"]
if sconf.keys():
print("\nChecking server connections:")
print("-" * 28)
excpts = {}
connect_timeout = 60
for s in servers:
if s in sconf:
try:
conn = Connection(sconf[s])
conn_check = Process(target=conn.connect)
conn_check.start()
counter = 0
while conn_check.is_alive():
time.sleep(1)
counter += 1
if counter > connect_timeout:
conn_check.kill()
raise Exception(
f"Connection was killed due to timeout ({connect_timeout}s)"
)
conn.release()
print(f"{s} connection: OK")
except Exception as e:
print(f"{s} connection: Error")
excpts[s] = e
if excpts:
print("\nExceptions:")
for k, v in excpts.items():
print(f"{k}: {v}")
class TornadoConsumer(object):
"""Non-blocking, Tornado ioloop based consumer"""
def __init__(self, broker_url, exchange, exchange_type, queue,
routing_key, durable=True, ssl=False, io_loop=None):
self.io_loop = io_loop or ioloop.IoLoop.instance()
self._conn = None
self._callbacks = []
self._broker_url = broker_url
self._ssl = ssl
self._exchange = Exchange(exchange, exchange_type, durable=durable)
self._queue = Queue(queue, exchange=self._exchange, routing_key=routing_key)
def add_callback(self, callback):
assert not self._conn
self._callbacks.append(callback)
def start(self):
self._conn = Connection(self._broker_url, self._ssl)
self._consumer = self._conn.Consumer(self._queue, callbacks=self._callbacks)
self.io_loop.add_handler(self._conn.fileno(), self._handle_event)
def stop(self):
self.io_loop.remove_handler(self._conn.fileno())
self._consumer.close()
self._conn.release()
def join(self, *args, **kwargs):
pass
def _handle_event(self):
self._conn.drain_nowait()
class AmqpBundle(object):
def __init__(self):
self.logger = logging.getLogger("amqp_bundle")
self.config_mapping = {
"amqp": {
"uri": None,
"autodisconnect": False,
"heartbeat": 0,
}
}
self.autodisconnect = None
self.event_listeners = [
(ConfigurationReadyEvent, self.configuration_ready),
(KernelShutdownEvent, self.kernel_shutdown),
]
self.connection = Connection()
self.injection_bindings = {Connection: self.connection}
def kernel_shutdown(self, event):
if self.autodisconnect:
self.logger.info("Disconnecting...")
self.connection.release()
self.logger.info("Disconnected")
def configuration_ready(self, event):
c = event.configuration.amqp
self.autodisconnect = c.autodisconnect
# First connect the topic manager to avoid lose messages
self.connection.__init__(c.uri, heartbeat=c.heartbeat)
self.connection.connect()
self.logger.info("Connected")
def _connect(self):
result = None
if not self.connected():
self.logger.info(
"Connect to RabbitMQ Broker %s" % (self.uri))
uri = self.uri
parsed_url = urlparse(uri)
uri = 'amqp://%s/%s' % (parsed_url.netloc, parsed_url.path)
result = Connection(uri)
try:
result.connect()
self.logger.info("Connected to AMQP Broker.")
except Exception as e:
result.release()
self.logger.error(
"Connection failure to %s: %s" % (uri, e))
else:
self.logger.debug("Allready connected")
return result
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)
self._connection.release()
class GenericPublisher(object):
"""Generic publisher class that specific publishers inherit from."""
def __init__(self, config, exchange=None, connect=True, **kwargs):
self.config = config
self.exchange = exchange
self.connection = None
if connect:
self.connect()
def connect(self):
if not self.connection:
self.connection = Connection(hostname=self.config.host,
port=self.config.port,
userid=self.config.user,
password=self.config.password,
virtual_host=self.config.vhost,
ssl=self.config.ssl)
def disconnect(self):
if self.connection:
self.connection.release()
self.connection = None
def publish(self, message):
"""Publishes a pulse message to the proper exchange."""
if not self.exchange:
raise InvalidExchange(self.exchange)
if not message:
raise MalformedMessage(message)
message._prepare()
if not self.connection:
self.connect()
producer = Producer(channel=self.connection,
exchange=Exchange(self.exchange, type='topic'),
routing_key=message.routing_key)
# The message is actually a simple envelope format with a payload and
# some metadata.
final_data = {}
final_data['payload'] = message.data
final_data['_meta'] = message.metadata.copy()
final_data['_meta'].update({
'exchange':
self.exchange,
'routing_key':
message.routing_key,
'serializer':
self.config.serializer,
'sent':
time_to_string(datetime.now(timezone(self.config.broker_timezone)))
})
producer.publish(final_data, serializer=self.config.serializer)
def inner(*args, **kwargs):
connection = Connection(settings.BROKER_URL)
connection.connect()
queue = connection.SimpleQueue(settings.QUEUE_NAME)
res = func(*args, **kwargs, queue=queue)
connection.release()
return res
def health_check_rabbitmq(broker_url):
try:
from kombu import Connection
conn = Connection(broker_url)
conn.connect()
conn.release()
return True, 'Rabbitmq is OK.', ''
except Exception as e:
return False, str(e), traceback.format_exc()
class GenericPublisher(object):
"""Generic publisher class that specific publishers inherit from."""
def __init__(self, config, exchange=None, connect=True, **kwargs):
self.config = config
self.exchange = exchange
self.connection = None
if connect:
self.connect()
def connect(self):
if not self.connection:
self.connection = Connection(hostname=self.config.host,
port=self.config.port,
userid=self.config.user,
password=self.config.password,
virtual_host=self.config.vhost,
ssl=self.config.ssl)
def disconnect(self):
if self.connection:
self.connection.release()
self.connection = None
def publish(self, message):
"""Publishes a pulse message to the proper exchange."""
if not self.exchange:
raise InvalidExchange(self.exchange)
if not message:
raise MalformedMessage(message)
message._prepare()
if not self.connection:
self.connect()
producer = Producer(channel=self.connection,
exchange=Exchange(self.exchange, type='topic'),
routing_key=message.routing_key)
# The message is actually a simple envelope format with a payload and
# some metadata.
final_data = {}
final_data['payload'] = message.data
final_data['_meta'] = message.metadata.copy()
final_data['_meta'].update({
'exchange': self.exchange,
'routing_key': message.routing_key,
'serializer': self.config.serializer,
'sent': time_to_string(datetime.now(timezone(self.config.broker_timezone)))
})
producer.publish(final_data, serializer=self.config.serializer)
def publish(queue, routing_key, body):
connection = Connection(app.config.get('AMQP_QUEUE_URI'))
try:
with connection.Producer(serializer='json') as producer:
producer.publish(body,
exchange=exchange,
routing_key=routing_key,
declare=[queue])
finally:
connection.release()
class InteractiveExchange:
'''
A context manager that provides the interactive exchange for publishing
messages.
'''
def __init__(self, task):
'''
Accessing an Interactive Exchange need a task only for the generic info in it,
and so this can be an instantiaton of Interactive()
:param task: A task that provides the exchange name and some logging formatters
'''
self.task = task
def __enter__(self):
# We have a task that we need only for the exchnage name and some logging formatters.
task = self.task
log.debug(f'Connecting to: {current_app.conf.broker_url}')
self.connection = Connection(current_app.conf.broker_url)
# Connection is lazy. Force a connection now.
self.connection.connect()
# Get the connection object for channel construction
#
# This syntax is odd, but we can blame kombu for that.
#
# self.connection is of type: kombu.connection.Connection
#
# and it has an attribute:
#
# self.connection.connection which is of type: kombu.transport.pyamqp.Connection
#
# So sname names, different things. The convoluted life of the
# word "connection". But we do need a kombu.transport.pyamqp.Connection
# object to create a channel. We can't do that on its parent.
c = self.connection.connection
log.debug(f'\tConnection: {task.connection_names(c)[1]}')
# Create a channel on the connection and log it in the RabbitMQ webmonitor format
ch = c.channel()
self.channel = ch
log.debug(f'\tChannel: {task.channel_names(ch)[1]}')
x = Exchange(task.exchange_name, channel=ch)
self.exchange = x
log.debug(f'\tExchange: {x.name}')
return x
def __exit__(self, type, value, traceback): # @ReservedAssignment
self.connection.release()
def publish(queue, routing_key, body):
connection = Connection(app.config.get('AMQP_QUEUE_URI'))
try:
with connection.Producer(serializer='json') as producer:
producer.publish(body,
exchange=exchange,
routing_key=routing_key,
declare=[queue])
finally:
connection.release()
def test_simple_queueing(self):
conn = Connection('sqlalchemy+sqlite:///:memory:')
conn.connect()
try:
channel = conn.channel()
assert channel.queue_cls.__table__.name == 'kombu_queue'
assert channel.message_cls.__table__.name == 'kombu_message'
channel._put('celery', 'DATA_SIMPLE_QUEUEING')
assert channel._get('celery') == 'DATA_SIMPLE_QUEUEING'
finally:
conn.release()
def listen(self):
logger.debug('AutophonePulseMonitor: start shared_lock.acquire')
connection = None
restart = True
while restart:
restart = False
self.shared_lock.acquire()
try:
# connection does not connect to the server until
# either the connection.connect() method is called
# explicitly or until kombu calls it implicitly as
# needed.
connection = Connection(hostname=self.hostname,
userid=self.userid,
password=self.password,
virtual_host=self.virtual_host,
port=DEFAULT_SSL_PORT,
ssl=True)
consumer = connection.Consumer(self.queues,
callbacks=[self.handle_message],
accept=['json'],
auto_declare=False)
for queue in self.queues:
queue(connection).queue_declare(passive=False)
queue(connection).queue_bind()
with consumer:
while not self._stopping.is_set():
try:
logger.debug('AutophonePulseMonitor shared_lock.release')
self.shared_lock.release()
connection.drain_events(timeout=self.timeout)
except socket.timeout:
pass
except KeyboardInterrupt:
raise
finally:
logger.debug('AutophonePulseMonitor shared_lock.acquire')
self.shared_lock.acquire()
logger.debug('AutophonePulseMonitor.listen: stopping')
except:
logger.exception('AutophonePulseMonitor Exception')
if connection:
connection.release()
restart = True
time.sleep(1)
finally:
logger.debug('AutophonePulseMonitor exit shared_lock.release')
if connection and not restart:
connection.release()
self.shared_lock.release()
def test_clone(self):
hostname = 'sqlite:///celerydb.sqlite'
x = Connection('+'.join(['sqla', hostname]))
try:
assert x.uri_prefix == 'sqla'
assert x.hostname == hostname
clone = x.clone()
try:
assert clone.hostname == hostname
assert clone.uri_prefix == 'sqla'
finally:
clone.release()
finally:
x.release()
class TriggerWatcher(ConsumerMixin):
TRIGGER_WATCH_Q = reactor.get_trigger_queue('st2.trigger.watch',
routing_key='#')
def __init__(self, create_handler, update_handler, delete_handler):
self._handlers = {
publishers.CREATE_RK: create_handler,
publishers.UPDATE_RK: update_handler,
publishers.DELETE_RK: delete_handler
}
self.connection = None
self._thread = None
def get_consumers(self, Consumer, channel):
return [Consumer(queues=[self.TRIGGER_WATCH_Q],
accept=['pickle'],
callbacks=[self.process_task])]
def process_task(self, body, message):
LOG.debug('process_task')
LOG.debug(' body: %s', body)
LOG.debug(' message.properties: %s', message.properties)
LOG.debug(' message.delivery_info: %s', message.delivery_info)
routing_key = message.delivery_info.get('routing_key', '')
handler = self._handlers.get(routing_key, None)
if not handler:
LOG.debug('Skipping message %s as no handler was found.')
return
try:
handler(body)
except:
LOG.exception('handling failed. Message body : %s', body)
finally:
message.ack()
def start(self):
try:
self.connection = Connection(cfg.CONF.messaging.url)
self._thread = eventlet.spawn(self.run)
except:
LOG.exception('Failed to start watcher.')
self.connection.release()
def stop(self):
try:
self._thread = eventlet.kill(self._thread)
finally:
self.connection.release()
def check_rabbit_mq():
conn = Connection(settings.AMQP_URL)
for attempt in range(1, 11):
timeout = attempt * 2
try:
conn.connect()
# Check connection
conn.release()
return True
except OSError as e:
logging.warning(
"Cannot connect to AMQP at {}. Retrying in {} secs...".format(
settings.AMQP_URL, timeout))
time.sleep(timeout)
return False
class Amqp(object):
def __init__(self, url, exchange, queue, routing_key):
self.conn = Connection(url)
self.exchange = Exchange(exchange, 'direct')
self.routing_key = routing_key
self.queue = Queue(queue, self.exchange, self.routing_key)
self.producer = None
self.consumer = None
def send(self, obj):
if not self.producer:
self.producer = self.conn.Producer()
self.producer.publish(obj, exchange=self.exchange,
routing_key=self.routing_key,
declare=[self.queue],
serializer='json', compression='zlib')
def poll(self, cb_func):
if not self.consumer:
self.consumer = self.conn.Consumer(self.queue,
callbacks=[cb_func])
self.consumer.qos(prefetch_count=1)
self.consumer.consume()
while True:
self.conn.drain_events()
def _release(self):
if self.consumer:
self.consumer.close()
self.consumer = None
if self.producer:
self.producer.close()
self.producer = None
if self.conn:
self.conn.release()
self.conn = None
def __enter__(self):
return self
def __exit__(self, exec_type, exc_value, traceback):
self._release()
class Amqp(object):
def __init__(self, url, exchange, queue, routing_key):
self.conn = Connection(url)
self.exchange = Exchange(exchange, 'direct')
self.routing_key = routing_key
self.queue = Queue(queue, self.exchange, self.routing_key)
self.producer = None
self.consumer = None
def send(self, obj):
if not self.producer:
self.producer = self.conn.Producer()
self.producer.publish(obj,
exchange=self.exchange,
routing_key=self.routing_key,
declare=[self.queue],
serializer='json',
compression='zlib')
def poll(self, cb_func):
if not self.consumer:
self.consumer = self.conn.Consumer(self.queue, callbacks=[cb_func])
self.consumer.qos(prefetch_count=1)
self.consumer.consume()
while True:
self.conn.drain_events()
def _release(self):
if self.consumer:
self.consumer.close()
self.consumer = None
if self.producer:
self.producer.close()
self.producer = None
if self.conn:
self.conn.release()
self.conn = None
def __enter__(self):
return self
def __exit__(self, exec_type, exc_value, traceback):
self._release()
def handle(self, command):
command_body = {
"command": command.command,
"argument": command.argument,
"user_id": command.user_id,
"chatroom_id": command.chatroom_id,
}
connection = Connection(settings.AMQP_ADDRESS)
connection.connect()
producer = connection.Producer()
producer.publish(
command_body,
retry=True,
exchange=settings.EXCHANGE,
routing_key=settings.SENDER_ROUTING_KEY,
declare=[settings.BOT_QUEUE],
)
connection.release()
def test_rabbitmq_on_startup(retry: int = 3, wait: float = 10):
"""
Tests connection to RabbitMQ on application start up.
:param retry: amount of connection retries
:param wait: seconds to wait until connection retry
"""
logging.info("Trying to connect to RabbitMQ ...")
rabbitmq_connection = Kombu(rabbitmq_url)
for i in range(0, retry):
try:
rabbitmq_connection.connect()
logging.info("Successfully connected to RabbitMQ!")
rabbitmq_connection.release()
return True
except KombuConnectionError or ConnectionRefusedError:
logging.warning("RabbitMQ is not yet ready!")
time.sleep(wait)
raise Exception("Cannot connect to RabbitMQ!")
class EventsPublisher(object):
"""Generic publisher class that specific publishers inherit from."""
def __init__(self, connection_url, routing_key=None):
self.connection_url = connection_url
self.connection = None
self.producer = None
self.default_routing_key = routing_key
self.exchange = Exchange("events", type="topic")
self.logger = logger
def connect(self):
self.connection = Connection(self.connection_url)
self.exchange = self.exchange(self.connection.channel)
def disconnect(self):
if self.connection:
self.connection.release()
self.connection = None
self.producer = None
def log(self, message):
self.logger.info("{0}".format(message))
def publish(self, message, routing_key=None):
if not self.connection:
self.connect()
if not self.producer:
self.producer = Producer(self.connection, self.exchange)
routing_key = routing_key or self.default_routing_key
if not routing_key:
raise Exception("Routing key not specified")
self.log("Publishing to exchange {0} with routing key {1}".format(
self.exchange, routing_key
))
self.producer.publish(message,
exchange=self.exchange,
routing_key=routing_key)
class EventsPublisher(object):
"""Generic publisher class that specific publishers inherit from."""
def __init__(self, connection_url, routing_key=None):
self.connection_url = connection_url
self.connection = None
self.producer = None
self.default_routing_key = routing_key
self.exchange = Exchange("events", type="topic")
self.logger = logger
def connect(self):
self.connection = Connection(self.connection_url)
self.exchange = self.exchange(self.connection.channel)
def disconnect(self):
if self.connection:
self.connection.release()
self.connection = None
self.producer = None
def log(self, message):
self.logger.info("{0}".format(message))
def publish(self, message, routing_key=None):
if not self.connection:
self.connect()
if not self.producer:
self.producer = Producer(self.connection, self.exchange)
routing_key = routing_key or self.default_routing_key
if not routing_key:
raise Exception("Routing key not specified")
self.log("Publishing to exchange {0} with routing key {1}".format(
self.exchange, routing_key
))
self.producer.publish(message,
exchange=self.exchange,
routing_key=routing_key)
def __call__(self):
"""
Broker URL must start with "amqp://" else return None
Return True if can establish connection or not.
"""
if self.broker_url.startswith("amqp://"):
try:
connection = Connection(self.broker_url, connection_timeout=1)
try:
connection.connect()
status = connection.connected
finally:
connection.release()
except Exception as exc:
logger.exception(exc)
status = False
return status
else:
return None
class SocketBrokerClient:
"""
Base class for web socket notification using broker (redis or rabbitmq)
"""
connection = None
def __init__(self, url, exchange_name=None):
self.url = url
self.connect()
self.exchange_name = exchange_name if exchange_name else celery_queue(
"socket_notification")
self.channel = self.connection.channel()
self.socket_exchange = Exchange(self.exchange_name,
type="fanout",
channel=self.channel)
self.socket_exchange.declare()
def open(self):
"""Test if connection is open.
True if connected else false
:return bool:
"""
return self.connection and self.connection.connected
def connect(self):
self._close()
logger.info("Connecting to broker {}".format(self.url))
self.connection = Connection(self.url, heartbeat=WS_HEART_BEAT)
self.connection.connect()
logger.info("Connected to broker {}".format(self.url))
def _close(self):
if hasattr(self, "connection") and self.connection:
logger.info("Closing connecting to broker {}".format(self.url))
self.connection.release()
self.connection = None
logger.info("Connection closed to broker {}".format(self.url))
def close(self):
self._close()
class SocketBrokerClient:
"""
Base class for web socket notification using broker (redis or rabbitmq)
"""
connection = None
def __init__(self, url):
self.url = url
self.connect()
self.channel = self.connection.channel()
self.socket_exchange = Exchange(exchange_name,
type='fanout',
channel=self.channel)
self.socket_exchange.declare()
def open(self):
"""Test if connection is open.
True if connected else false
:return bool:
"""
return self.connection and self.connection.connected
def connect(self):
self._close()
logger.info('Connecting to broker {}'.format(self.url))
self.connection = Connection(self.url)
self.connection.connect()
logger.info('Connected to broker {}'.format(self.url))
def _close(self):
if hasattr(self, 'connection') and self.connection:
logger.info('Closing connecting to broker {}'.format(self.url))
self.connection.release()
self.connection = None
logger.info('Connection closed to broker {}'.format(self.url))
def close(self):
self._close()
class SocketBrokerClient:
"""
Base class for web socket notification using broker (redis or rabbitmq)
"""
connection = None
def __init__(self, url, exchange_name=None):
self.url = url
self.connect()
self.exchange_name = exchange_name if exchange_name else celery_queue('socket_notification')
self.channel = self.connection.channel()
self.socket_exchange = Exchange(self.exchange_name, type='fanout', channel=self.channel)
self.socket_exchange.declare()
def open(self):
"""Test if connection is open.
True if connected else false
:return bool:
"""
return self.connection and self.connection.connected
def connect(self):
self._close()
logger.info('Connecting to broker {}'.format(self.url))
self.connection = Connection(self.url, heartbeat=WS_HEART_BEAT)
self.connection.connect()
logger.info('Connected to broker {}'.format(self.url))
def _close(self):
if hasattr(self, 'connection') and self.connection:
logger.info('Closing connecting to broker {}'.format(self.url))
self.connection.release()
self.connection = None
logger.info('Connection closed to broker {}'.format(self.url))
def close(self):
self._close()
def _examine_connection(s, sconf, excpts):
connect_timeout = 60
try:
conn = Connection(sconf[s])
conn_check = ConnProcess(target=conn.connect)
conn_check.start()
counter = 0
while conn_check.is_alive():
time.sleep(1)
counter += 1
if counter > connect_timeout:
conn_check.kill()
raise Exception(f"Connection was killed due to timeout ({connect_timeout}s)")
conn.release()
if conn_check.exception:
error, traceback = conn_check.exception
raise error
except Exception as e:
print(f"{s} connection: Error")
excpts[s] = e
else:
print(f"{s} connection: OK")
def _rabbit_mq(messages):
conn = Connection(host)
channel = conn.channel()
exchange = Exchange(exchange_name, type="direct")
queue = Queue(name=queue_name,
exchange=exchange,
routing_key=routing_key)
queue.maybe_bind(conn)
queue.declare()
producer = Producer(exchange=exchange,
channel=channel,
routing_key=routing_key)
for message in messages:
# as_dict = message.asDict(recursive=True)
producer.publish(message)
channel.close()
conn.release()
return messages
class SocketBrokerClient:
"""
Base class for web socket notification using broker (redis or rabbitmq)
"""
connection = None
def __init__(self, url):
self.url = url
self.connect()
self.channel = self.connection.channel()
self.socket_exchange = Exchange(exchange_name, type="fanout", channel=self.channel)
self.socket_exchange.declare()
def open(self):
"""
True if connected else false
:return bool:
"""
return self.connection and self.connection.connected
def connect(self):
self._close()
logger.info("Connecting to broker {}".format(self.url))
self.connection = Connection(self.url)
self.connection.connect()
logger.info("Connected to broker {}".format(self.url))
def _close(self):
if hasattr(self, "connection") and self.connection:
logger.info("Closing connecting to broker {}".format(self.url))
self.connection.release()
self.connection = None
logger.info("Connection closed to broker {}".format(self.url))
def close(self):
self._close()
def check_server_access(sconf):
servers = ["broker server", "results server"]
if sconf.keys():
print("\nChecking server connections:")
print("-" * 28)
excpts = {}
for s in servers:
if s in sconf:
try:
conn = Connection(sconf[s])
conn.connect()
conn.release()
print(f"{s} connection: OK")
except Exception as e:
print(f"{s} connection: Error")
excpts[s] = e
if excpts:
print("\nExceptions:")
for k, v in excpts.items():
print(f"{k}: {v}")
def send_event(self, event):
"""
Send event to AMQP bus.
@param event AMQP event to send.
@type event dict
"""
from kombu import Connection
from kombu.pools import producers
conn = Connection(hostname=self.host, userid=self.user, virtual_host=self.vhost)
producer = producers[conn].acquire(block=True)
producer.publish(
event,
serializer='json',
exchange=self.exchange,
routing_key=self.event_get_rk(event)
)
producer.release()
conn.release()
def handle(self, *args, **options):
routing_key = options["routing_key"]
connection_url = options["connection_url"]
userid = urlparse(connection_url).username
payload_file = options["payload_file"]
exchange_name = "exchange/{}/jobs".format(userid)
connection = Connection(connection_url)
exchange = Exchange(exchange_name, type="topic")
producer = Producer(connection,
exchange,
routing_key=routing_key,
auto_declare=True)
self.stdout.write("Published to exchange: {}".format(exchange_name))
with open(payload_file) as f:
body = f.read()
try:
producer.publish(body)
finally:
connection.release()
def revoke(app_name, app_env, task_uuid: Union[str, List[str]], args):
# Check if agent is local
if not settings.LEEK_ENABLE_AGENT:
return responses.control_operations_not_supported
# Retrieve subscription
found, subscription = lookup_subscription(app_name, app_env)
if not found:
return responses.task_retry_subscription_not_found
# Prepare connection/producer
# noinspection PyBroadException
try:
connection = Connection(subscription["broker"])
connection.ensure_connection(max_retries=2)
producer = connection.Producer()
except AccessRefused:
return responses.wrong_access_refused
except Exception:
return responses.broker_not_reachable
arguments = {
"task_id": task_uuid,
**args,
}
# noinspection PyBroadException
try:
broadcast_worker_command("revoke", arguments, producer)
except Exception as ex:
logger.error(ex)
return responses.task_revocation_failed
connection.release()
revocation_count = len(task_uuid) if isinstance(task_uuid, List) else 1
return {"acknowledged": True, "revocation_count": revocation_count}, 200
def handle(self, *args, **options):
routing_key = options["routing_key"]
connection_url = options["connection_url"]
userid = urlparse(connection_url).username
payload_file = options["payload_file"]
exchange_name = "exchange/{}/jobs".format(userid)
connection = Connection(connection_url)
exchange = Exchange(exchange_name, type="topic")
producer = Producer(connection,
exchange,
routing_key=routing_key,
auto_declare=True)
self.stdout.write("Published to exchange: {}".format(exchange_name))
with open(payload_file) as f:
body = f.read()
try:
producer.publish(body)
finally:
connection.release()
class Publisher(object):
"""
Mimic GenericPublisher https://github.com/bhearsum/mozillapulse/blob/master/mozillapulse/publishers.py
"""
@use_settings
def __init__(
self,
exchange, # name of the Pulse exchange
host='pulse.mozilla.org', # url to connect,
port=5671, # tcp port
user=None,
password=None,
vhost="/",
start=0, # USED AS STARTING POINT FOR ASSIGNING THE _meta.count ATTRIBUTE
ssl=True,
applabel=None,
heartbeat=False, # True to also get the Pulse heartbeat message
durable=False, # True to keep queue after shutdown
serializer='json',
broker_timezone='GMT',
settings=None
):
self.settings = settings
self.connection = None
self.count = 0
def connect(self):
if not self.connection:
self.connection = Connection(
hostname=self.settings.host,
port=self.settings.port,
userid=self.settings.user,
password=self.settings.password,
virtual_host=self.settings.vhost,
ssl=self.settings.ssl
)
def disconnect(self):
if self.connection:
self.connection.release()
self.connection = None
def send(self, topic, message):
"""Publishes a pulse message to the proper exchange."""
if not message:
Log.error("Expecting a message")
message._prepare()
if not self.connection:
self.connect()
producer = Producer(
channel=self.connection,
exchange=Exchange(self.settings.exchange, type='topic'),
routing_key=topic
)
# The message is actually a simple envelope format with a payload and
# some metadata.
final_data = Dict(
payload=message.data,
_meta=set_default({
'exchange': self.settings.exchange,
'routing_key': message.routing_key,
'serializer': self.settings.serializer,
'sent': time_to_string(datetime.datetime.now(timezone(self.settings.broker_timezone))),
'count': self.count
}, message.metadata)
)
producer.publish(jsons.scrub(final_data), serializer=self.settings.serializer)
self.count += 1
class RabbitMQSource(StoqSourcePlugin):
def __init__(self):
super().__init__()
def activate(self, stoq):
self.stoq = stoq
super().activate()
self.amqp_publish_conn = None
def ingest(self):
"""
Monitor AMQP for messages
"""
if self.stoq.worker.name:
# Define our RabbitMQ route
routing_key = self.stoq.worker.name
# If this is an error message, let's make sure our queue
# has "-errors" affixed to it
if self.stoq.worker.error_queue is True:
routing_key = routing_key + "-errors".strip()
exchange = Exchange(self.exchange_name, type=self.exchange_type)
queue = Queue(routing_key, exchange, routing_key=routing_key)
# Setup our broker connection with RabbitMQ
with Connection(hostname=self.host,
port=self.port,
userid=self.user,
password=self.password,
virtual_host=self.virtual_host) as conn:
consumer = Consumer(conn, queue,
callbacks=[self.queue_callback])
consumer.qos(prefetch_count=int(self.prefetch))
consumer.consume()
while True:
conn.drain_events()
else:
self.stoq.log.error("No worker name defined!")
def queue_callback(self, amqp_message_data, amqp_message_handler):
try:
# Setup the amqp message for parsing
kwargs = self.stoq.loads(amqp_message_data)
# Send the message to the worker
self.stoq.worker.start(**kwargs)
except Exception as e:
# Something went wrong, let's publish to the error queue. and
# append the error message.
kwargs['err'] = str(e)
self.stoq.log.error(kwargs)
self.publish_connect()
self.publish(kwargs, self.stoq.worker.name, err=True)
self.publish_release()
# Acknowledge our message in amqp so we can get the next one.
amqp_message_handler.ack()
def publish_connect(self):
"""
Connect to AMQP to publish a message
:returns: AMQP connection object for publishing
:rtype: kombu.Connection object
"""
self.amqp_exchange = Exchange(self.exchange_name,
type=self.exchange_type)
self.amqp_publish_conn = Connection(hostname=self.host,
port=self.port,
userid=self.user,
password=self.password,
virtual_host=self.virtual_host)
return self.amqp_publish_conn
def publish_release(self):
"""
Release AMQP connection used for publishing
"""
return self.amqp_publish_conn.release()
def publish(self, msg, routing_key, err=False):
"""
Publish a message to AMQP
:param dict msg: Message to be published, preferrably json
:param bytes routing_key: Route to be used, should be name of worker
:param err err: Define whether we should process error queue
"""
# Make sure we have a valid connection to RabbitMQ
if not self.amqp_publish_conn:
self.publish_connect()
# If this is an error message, let's make sure our queue
# has "-errors" affixed to it
if err:
routing_key = routing_key + "-errors".strip()
# Define the queue so we can ensure it is declared before
# publishing to it
queue = Queue(routing_key, self.amqp_exchange, routing_key=routing_key)
with producers[self.amqp_publish_conn].acquire(block=True) as producer:
producer.publish(self.stoq.dumps(msg),
exchange=self.amqp_exchange,
declare=[self.amqp_exchange, queue],
routing_key=routing_key)
def __errback(self, exc, interval):
"""
Error handling for AMQP publishing
"""
self.stoq.log.warn("Unable to publish message: {}. Retry in {}s.".format(
exc, interval))
def kombu_connection_test():
connection = Connection('amqp://*****:*****@121.41.4.137:5672//')
connection.connect()
print connection.connected
connection.release()
class SensorWatcher(ConsumerMixin):
def __init__(self,
create_handler,
update_handler,
delete_handler,
queue_suffix=None):
"""
:param create_handler: Function which is called on SensorDB create event.
:type create_handler: ``callable``
:param update_handler: Function which is called on SensorDB update event.
:type update_handler: ``callable``
:param delete_handler: Function which is called on SensorDB delete event.
:type delete_handler: ``callable``
"""
# TODO: Handle sensor type filtering using routing key
self._create_handler = create_handler
self._update_handler = update_handler
self._delete_handler = delete_handler
self._sensor_watcher_q = self._get_queue(queue_suffix)
self.connection = None
self._updates_thread = None
self._handlers = {
publishers.CREATE_RK: create_handler,
publishers.UPDATE_RK: update_handler,
publishers.DELETE_RK: delete_handler
}
def get_consumers(self, Consumer, channel):
consumers = [
Consumer(queues=[self._sensor_watcher_q],
accept=['pickle'],
callbacks=[self.process_task])
]
return consumers
def process_task(self, body, message):
LOG.debug('process_task')
LOG.debug(' body: %s', body)
LOG.debug(' message.properties: %s', message.properties)
LOG.debug(' message.delivery_info: %s', message.delivery_info)
routing_key = message.delivery_info.get('routing_key', '')
handler = self._handlers.get(routing_key, None)
try:
if not handler:
LOG.info('Skipping message %s as no handler was found.',
message)
return
try:
handler(body)
except Exception as e:
LOG.exception(
'Handling failed. Message body: %s. Exception: %s', body,
e.message)
finally:
message.ack()
def start(self):
try:
self.connection = Connection(transport_utils.get_messaging_urls())
self._updates_thread = eventlet.spawn(self.run)
except:
LOG.exception('Failed to start sensor_watcher.')
self.connection.release()
def stop(self):
LOG.debug('Shutting down sensor watcher.')
try:
if self._updates_thread:
self._updates_thread = eventlet.kill(self._updates_thread)
if self.connection:
channel = self.connection.channel()
bound_sensor_watch_q = self._sensor_watcher_q(channel)
try:
bound_sensor_watch_q.delete()
except:
LOG.error('Unable to delete sensor watcher queue: %s',
self._sensor_watcher_q)
finally:
if self.connection:
self.connection.release()
@staticmethod
def _get_queue(queue_suffix):
queue_name = queue_utils.get_queue_name(
queue_name_base='st2.sensor.watch',
queue_name_suffix=queue_suffix,
add_random_uuid_to_suffix=True)
return reactor.get_sensor_cud_queue(queue_name, routing_key='#')
class SensorWatcher(ConsumerMixin):
def __init__(self, create_handler, update_handler, delete_handler,
queue_suffix=None):
"""
:param create_handler: Function which is called on SensorDB create event.
:type create_handler: ``callable``
:param update_handler: Function which is called on SensorDB update event.
:type update_handler: ``callable``
:param delete_handler: Function which is called on SensorDB delete event.
:type delete_handler: ``callable``
"""
# TODO: Handle sensor type filtering using routing key
self._create_handler = create_handler
self._update_handler = update_handler
self._delete_handler = delete_handler
self._sensor_watcher_q = self._get_queue(queue_suffix)
self.connection = None
self._updates_thread = None
self._handlers = {
publishers.CREATE_RK: create_handler,
publishers.UPDATE_RK: update_handler,
publishers.DELETE_RK: delete_handler
}
def get_consumers(self, Consumer, channel):
consumers = [Consumer(queues=[self._sensor_watcher_q],
accept=['pickle'],
callbacks=[self.process_task])]
return consumers
def process_task(self, body, message):
LOG.debug('process_task')
LOG.debug(' body: %s', body)
LOG.debug(' message.properties: %s', message.properties)
LOG.debug(' message.delivery_info: %s', message.delivery_info)
routing_key = message.delivery_info.get('routing_key', '')
handler = self._handlers.get(routing_key, None)
try:
if not handler:
LOG.info('Skipping message %s as no handler was found.', message)
return
try:
handler(body)
except Exception as e:
LOG.exception('Handling failed. Message body: %s. Exception: %s',
body, e.message)
finally:
message.ack()
def start(self):
try:
self.connection = Connection(transport_utils.get_messaging_urls())
self._updates_thread = eventlet.spawn(self.run)
except:
LOG.exception('Failed to start sensor_watcher.')
self.connection.release()
def stop(self):
try:
if self._updates_thread:
self._updates_thread = eventlet.kill(self._updates_thread)
finally:
if self.connection:
self.connection.release()
@staticmethod
def _get_queue(queue_suffix):
queue_name = queue_utils.get_queue_name(queue_name_base='st2.sensor.watch',
queue_name_suffix=queue_suffix,
add_random_uuid_to_suffix=True
)
return reactor.get_sensor_cud_queue(queue_name, routing_key='#')
class Amqp(Thread):
def __init__(self,
host="localhost",
port=5672,
userid="guest",
password="******",
virtual_host="canopsis",
exchange_name="canopsis",
logging_name="Amqp",
logging_level=INFO,
read_config_file=True,
auto_connect=True,
on_ready=None,
max_retries=5):
super(Amqp, self).__init__()
self.logger = getLogger(logging_name)
self.host = host
self.port = port
self.userid = userid
self.password = password
self.virtual_host = virtual_host
self.exchange_name = exchange_name
self.logging_level = logging_level
if read_config_file:
self.read_config("amqp")
self.amqp_uri = "amqp://{0}:{1}@{2}:{3}/{4}".format(
self.userid, self.password, self.host, self.port,
self.virtual_host)
self.logger.setLevel(logging_level)
# Event sent try count before event drop in case of connection problem
if max_retries != 5:
self.logger.info('Custom retries value : {} {}'.format(
max_retries, type(max_retries)))
self.max_retries = max_retries
self.exchange_name_events = exchange_name + ".events"
self.exchange_name_alerts = exchange_name + ".alerts"
self.exchange_name_incidents = exchange_name + ".incidents"
self.chan = None
self.conn = None
self.connected = False
self.on_ready = on_ready
self.RUN = True
self.exchanges = {}
self.queues = {}
self.producers = None
self.paused = False
self.connection_errors = (ConnectionError, error, IOError, OSError)
# Create exchange
self.logger.debug("Create exchanges object")
for exchange_name in [
self.exchange_name, self.exchange_name_events,
self.exchange_name_alerts, self.exchange_name_incidents
]:
self.logger.debug(' + {}'.format(exchange_name))
self.get_exchange(exchange_name)
if auto_connect:
self.connect()
self.logger.debug("Object canamqp initialized")
def run(self):
"""
main thread loop
- connects to rabbit
- calls drain_events
- disconnects on error
"""
self.logger.debug("Start thread ...")
reconnect = False
while self.RUN:
self.connect()
if self.connected:
self.init_queue(reconnect=reconnect)
self.logger.debug("Drain events ...")
while self.RUN:
try:
if not self.paused:
self.conn.drain_events(timeout=0.5)
else:
sleep(0.5)
except KombuSerializationError as exc:
self.logger.error(
"Kombu serialization error: invalid message "
"received: {}".format(exc))
except timeout:
pass
except self.connection_errors as err:
self.logger.error(
"Connection error ! ({})".format(err))
break
except KombuSerializationError as exc:
self.logger.error(
"Kombu serialization error: invalid message "
"received: {}".format(exc))
except Exception as err:
self.logger.error("Error: {} ({})".format(
err, type(err)))
print_exc(file=stdout)
break
self.disconnect()
if self.RUN:
self.logger.error(
'Connection lost, try to reconnect in few seconds ...')
reconnect = True
self.wait_connection(timeout=5)
self.logger.debug("End of thread ...")
def stop(self):
self.logger.debug("Stop thread ...")
self.RUN = False
def connect(self):
"""
Create the connection
Inits the producers
Init the only channel
Declares exchanges
"""
if not self.connected:
self.logger.info("Connect to AMQP Broker (%s:%s)" %
(self.host, self.port))
self.conn = Connection(self.amqp_uri)
try:
self.logger.debug(" + Connect")
self.conn.connect()
self.logger.info("Connected to AMQP Broker.")
self.producers = pools.Producers(limit=10)
self.connected = True
except Exception as err:
self.conn.release()
self.logger.error("Impossible to connect ({})".format(err))
if self.connected:
self.logger.debug(" + Open channel")
try:
self.chan = self.conn.channel()
self.logger.info("Channel openned. Ready to send messages")
try:
# Declare exchange
self.logger.debug("Declare exchanges")
for exchange_name in self.exchanges:
self.logger.debug(" + {}".format(exchange_name))
self.exchanges[exchange_name](self.chan).declare()
except Exception as err:
self.logger.error(
"Impossible to declare exchange ({})".format(err))
except Exception as err:
self.logger.error(err)
else:
self.logger.debug("Already connected")
def get_exchange(self, name):
"""
Returns an exchange if stored in self.exchanges.
Otherwise, creates it and returns it
:param string name: name of the exchange to get/create
:rtype: Exchange|None
"""
if name:
try:
return self.exchanges[name]
except Exception:
if name == "amq.direct":
self.exchanges[name] = Exchange(name,
"direct",
durable=True)
else:
self.exchanges[name] = Exchange(name,
"topic",
durable=True,
auto_delete=False)
return self.exchanges[name]
else:
return None
def init_queue(self, reconnect=False):
"""
Loads queue settings
Creates queues and attaches the same channel to each of them
Binds queues to exchange
Revives or creates consumers
calls consume
"""
if self.queues:
self.logger.debug("Init queues")
for queue_name in self.queues.keys():
self.logger.debug(" + {}".format(queue_name))
qsettings = self.queues[queue_name]
if not qsettings['queue']:
self.logger.debug(" + Create queue")
# copy list
routing_keys = list(qsettings['routing_keys'])
routing_key = None
if len(routing_keys):
routing_key = routing_keys[0]
routing_keys = routing_keys[1:]
exchange = self.get_exchange(qsettings['exchange_name'])
if qsettings['exchange_name'] == "amq.direct" \
and not routing_key:
routing_key = queue_name
self.logger.debug("exchange: '{}', exclusive: {},"
" auto_delete: {},no_ack: {}".format(
qsettings['exchange_name'],
qsettings['exclusive'],
qsettings['auto_delete'],
qsettings['no_ack']))
qsettings['queue'] = Queue(
queue_name,
exchange=exchange,
routing_key=routing_key,
exclusive=qsettings['exclusive'],
auto_delete=qsettings['auto_delete'],
no_ack=qsettings['no_ack'],
channel=self.conn.channel())
qsettings['queue'].declare()
if len(routing_keys):
self.logger.debug(" + Bind on all routing keys")
for routing_key in routing_keys:
self.logger.debug(
" + routing_key: '{}'".format(routing_key))
try:
qsettings['queue'].bind_to(
exchange=exchange, routing_key=routing_key)
except Exception:
self.logger.error(
"You need upgrade your Kombu version ({})".
format(__version__))
if qsettings['consumer'] and not reconnect:
qsettings['consumer'].revive(self.chan)
elif not qsettings['consumer'] or reconnect:
self.logger.debug(" + Create Consumer")
qsettings['consumer'] = self.conn.Consumer(
qsettings['queue'], callbacks=[qsettings['callback']])
self.logger.debug(" + Consume queue")
qsettings['consumer'].consume()
if self.on_ready:
self.on_ready()
else:
self.logger.info('Queue already inited')
def add_queue(self,
queue_name,
routing_keys,
callback,
exchange_name=None,
no_ack=True,
exclusive=False,
auto_delete=True):
"""
Initializes the queue configuration
maps the callback on the queue
"""
c_routing_keys = []
if not isinstance(routing_keys, list):
if isinstance(routing_keys, basestring):
c_routing_keys = [routing_keys]
else:
c_routing_keys = routing_keys
# aka if rk is nor a list nor a basetring, leave it empty
if not exchange_name:
exchange_name = self.exchange_name
self.queues[queue_name] = {
'queue': False,
'consumer': False,
'queue_name': queue_name,
'routing_keys': c_routing_keys,
'callback': callback,
'exchange_name': exchange_name,
'no_ack': no_ack,
'exclusive': exclusive,
'auto_delete': auto_delete
}
def publish(self,
msg,
routing_key,
exchange_name=None,
serializer="json",
compression=None,
content_type=None,
content_encoding=None):
"""
Tries to publish an event
In case of failure, tries to reconnect and retry until (self.max_retries)
:returns: operation success status
:rtype: bool
"""
self.logger.warning("Publishing from old.rabbitmq.Amqp is deprecated")
operation_success = False
retries = 0
while not operation_success and retries < self.max_retries:
retries += 1
if self.connected:
if not exchange_name:
exchange_name = self.exchange_name
with self.producers[self.conn].acquire(block=True) as producer:
try:
_msg = msg.copy()
Amqp._clean_msg_for_serialization(_msg)
exchange = self.get_exchange(
exchange_name.encode('utf-8'))
producer.publish(_msg,
serializer=serializer,
compression=compression,
routing_key=routing_key,
exchange=exchange)
self.logger.debug('publish {} in exchange {}'.format(
routing_key, exchange_name))
operation_success = True
except AmqpStructError:
self.logger.warning(
'Malformed message: routing key is too long. '
'Cancelling message')
return False
except Exception:
self.logger.error(' + Impossible to send {}'.format(
traceback.format_exc()))
self.disconnect()
self.connect()
self.init_queue(reconnect=False)
else:
self.logger.error('Not connected ... try reconnecting')
self.connect()
if not operation_success:
# Event and it's information are buffered until next send retry
self.logger.info('Retry count {}'.format(retries))
if not operation_success:
# Event and it's information are buffered until next send retry
self.logger.error(
'Too much retries for event {}, give up'.format(routing_key))
return operation_success
@staticmethod
def _clean_msg_for_serialization(msg):
from bson import objectid
for key in msg:
if isinstance(msg[key], objectid.ObjectId):
msg[key] = str(msg[key])
def cancel_queues(self):
if self.connected:
for queue_name in self.queues.keys():
if self.queues[queue_name]['consumer']:
self.logger.debug(
" + Cancel consumer on {}".format(queue_name))
try:
self.queues[queue_name]['consumer'].cancel()
except Exception:
pass
del (self.queues[queue_name]['consumer'])
self.queues[queue_name]['consumer'] = False
del (self.queues[queue_name]['queue'])
self.queues[queue_name]['queue'] = False
def disconnect(self):
if self.connected:
self.logger.info("Disconnect from AMQP Broker")
self.cancel_queues()
# Force producers closing to permit a clean reconnect after
# ... especially on timeout errors
self.producers[self.conn].force_close_all()
for exchange in self.exchanges:
del exchange
self.exchanges = {}
try:
pools.reset()
except Exception as err:
self.logger.error(
"Impossible to reset kombu pools: {} ({})".format(
err, type(err)))
try:
self.conn.release()
del self.conn
except Exception as err:
self.logger.error(
"Impossible to release connection: {} ({})".format(
err, type(err)))
self.connected = False
def wait_connection(self, timeout=5):
i = 0
while self.RUN and not self.connected and i < (timeout * 2):
sleep(0.5)
i += 1
def read_config(self, name):
filename = join(root_path, 'etc', '{0}.conf'.format(name))
import ConfigParser
self.config = ConfigParser.RawConfigParser()
try:
self.config.read(filename)
section = 'master'
self.host = self.config.get(section, "host")
self.port = self.config.getint(section, "port")
self.userid = self.config.get(section, "userid")
self.password = self.config.get(section, "password")
self.virtual_host = self.config.get(section, "virtual_host")
self.exchange_name = self.config.get(section, "exchange_name")
except Exception as err:
self.logger.error(
"Can't to load configurations ({}), use default ...".format(
err))
def __del__(self):
self.stop()
class Publisher(object):
"""
Mimic GenericPublisher https://github.com/bhearsum/mozillapulse/blob/master/mozillapulse/publishers.py
"""
@override
def __init__(
self,
exchange, # name of the Pulse exchange
host='pulse.mozilla.org', # url to connect,
port=5671, # tcp port
user=None,
password=None,
vhost="/",
start=0, # USED AS STARTING POINT FOR ASSIGNING THE _meta.count ATTRIBUTE
ssl=True,
applabel=None,
heartbeat=False, # True to also get the Pulse heartbeat message
durable=False, # True to keep queue after shutdown
serializer='json',
broker_timezone='GMT',
kwargs=None):
self.settings = kwargs
self.connection = None
self.count = 0
def connect(self):
if not self.connection:
self.connection = Connection(hostname=self.settings.host,
port=self.settings.port,
userid=self.settings.user,
password=self.settings.password,
virtual_host=self.settings.vhost,
ssl=self.settings.ssl)
def disconnect(self):
if self.connection:
self.connection.release()
self.connection = None
def send(self, topic, message):
"""Publishes a pulse message to the proper exchange."""
if not message:
Log.error("Expecting a message")
message._prepare()
if not self.connection:
self.connect()
producer = Producer(channel=self.connection,
exchange=Exchange(self.settings.exchange,
type='topic'),
routing_key=topic)
# The message is actually a simple envelope format with a payload and
# some metadata.
final_data = Data(
payload=message.data,
_meta=set_default(
{
'exchange':
self.settings.exchange,
'routing_key':
message.routing_key,
'serializer':
self.settings.serializer,
'sent':
time_to_string(
datetime.datetime.now(
timezone(self.settings.broker_timezone))),
'count':
self.count
}, message.metadata))
producer.publish(jsons.scrub(final_data),
serializer=self.settings.serializer)
self.count += 1
class TriggerWatcher(ConsumerMixin):
sleep_interval = 0 # sleep to co-operatively yield after processing each message
def __init__(self, create_handler, update_handler, delete_handler,
trigger_types=None, queue_suffix=None, exclusive=False):
"""
:param create_handler: Function which is called on TriggerDB create event.
:type create_handler: ``callable``
:param update_handler: Function which is called on TriggerDB update event.
:type update_handler: ``callable``
:param delete_handler: Function which is called on TriggerDB delete event.
:type delete_handler: ``callable``
:param trigger_types: If provided, handler function will only be called
if the trigger in the message payload is included
in this list.
:type trigger_types: ``list``
:param exclusive: If the Q is exclusive to a specific connection which is then
single connection created by TriggerWatcher. When the connection
breaks the Q is removed by the message broker.
:type exclusive: ``bool``
"""
# TODO: Handle trigger type filtering using routing key
self._create_handler = create_handler
self._update_handler = update_handler
self._delete_handler = delete_handler
self._trigger_types = trigger_types
self._trigger_watch_q = self._get_queue(queue_suffix, exclusive=exclusive)
self.connection = None
self._load_thread = None
self._updates_thread = None
self._handlers = {
publishers.CREATE_RK: create_handler,
publishers.UPDATE_RK: update_handler,
publishers.DELETE_RK: delete_handler
}
def get_consumers(self, Consumer, channel):
return [Consumer(queues=[self._trigger_watch_q],
accept=['pickle'],
callbacks=[self.process_task])]
def process_task(self, body, message):
LOG.debug('process_task')
LOG.debug(' body: %s', body)
LOG.debug(' message.properties: %s', message.properties)
LOG.debug(' message.delivery_info: %s', message.delivery_info)
routing_key = message.delivery_info.get('routing_key', '')
handler = self._handlers.get(routing_key, None)
try:
if not handler:
LOG.debug('Skipping message %s as no handler was found.', message)
return
trigger_type = getattr(body, 'type', None)
if self._trigger_types and trigger_type not in self._trigger_types:
LOG.debug('Skipping message %s since\'t trigger_type doesn\'t match (type=%s)',
message, trigger_type)
return
try:
handler(body)
except Exception as e:
LOG.exception('Handling failed. Message body: %s. Exception: %s',
body, e.message)
finally:
message.ack()
eventlet.sleep(self.sleep_interval)
def start(self):
try:
self.connection = Connection(transport_utils.get_messaging_urls())
self._updates_thread = eventlet.spawn(self.run)
self._load_thread = eventlet.spawn(self._load_triggers_from_db)
except:
LOG.exception('Failed to start watcher.')
self.connection.release()
def stop(self):
try:
self._updates_thread = eventlet.kill(self._updates_thread)
self._load_thread = eventlet.kill(self._load_thread)
finally:
self.connection.release()
# Note: We sleep after we consume a message so we give a chance to other
# green threads to run. If we don't do that, ConsumerMixin will block on
# waiting for a message on the queue.
def on_consume_end(self, connection, channel):
super(TriggerWatcher, self).on_consume_end(connection=connection,
channel=channel)
eventlet.sleep(seconds=self.sleep_interval)
def on_iteration(self):
super(TriggerWatcher, self).on_iteration()
eventlet.sleep(seconds=self.sleep_interval)
def _load_triggers_from_db(self):
for trigger_type in self._trigger_types:
for trigger in Trigger.query(type=trigger_type):
LOG.debug('Found existing trigger: %s in db.' % trigger)
self._handlers[publishers.CREATE_RK](trigger)
@staticmethod
def _get_queue(queue_suffix, exclusive):
queue_name = queue_utils.get_queue_name(queue_name_base='st2.trigger.watch',
queue_name_suffix=queue_suffix,
add_random_uuid_to_suffix=True
)
return reactor.get_trigger_cud_queue(queue_name, routing_key='#', exclusive=exclusive)
class RpcClient(object):
"""Process a RPC queue and fetch the response."""
reply_received = False
messages = {}
corr_id_server_queue = {}
def __init__(self,
exchange=default_exchange,
prefix=None,
client_queue=None):
"""Constructor for client object.
:exchange: Exchange to use
:prefix: for automatic server routing key generation
:client_queue: name for the client queue. Default is <command>.client
"""
self._prefix = prefix
self._exchange = exchange
self.reply = None
self._client_queue = client_queue
self._queue = None
self._conn = Connection(**CONN_DICT)
def retrieve_messages(self):
"""Process the message queue and ack the one that matches the
correlation_id sent to the server.
:correlation_id: Expected correlation id (uuid).
:queue: The queue to read.
:returns: JSON object retrieved from the queue.
"""
logger.debug("Client queue: {!r}".format(self._client_queue))
client_queue = self._queue
logger.debug("connection is {!r}"
"is connected: {!r}".format(self._conn,
self._conn.connected))
for i in collect_replies(self._conn,
self._conn.channel(),
client_queue,
timeout=1,
limit=1,
callbacks=[self.ack_message]):
logger.debug("Received {!r}".format(i))
if self.reply is not None:
response = self.reply
self.reply = None
try:
client_queue.purge()
client_queue.delete()
self._conn.release()
except Exception:
logger.warn("Unable to purge and delete queues",
exc_info=True)
return response
def ack_message(self, body, message):
logger.info("Processing message: {!r} with body {!r}".format(message,
body))
if 'correlation_id' in message.properties:
corr_id = message.properties['correlation_id']
try:
self.messages.pop(corr_id)
server_queue = self.corr_id_server_queue.pop(corr_id)
logger.info(
"STOPREQUEST:%s;CORRELATION_ID:%s" %
(server_queue, corr_id))
self.reply = body
message.ack()
except KeyError as e:
logger.exception(e)
logger.error("Malformed message: {!r}".format(body))
def _setup_payload(self, command_name, data):
"""Setup the datastructure for either hase-like or standard.
:command_name: the name of the command
:data: data to be sent with the command
:returns: payload for the request
"""
return data
def _prepare_client_queue(self, command_name):
"""Setup a client queue based on the command.
"""
if self._client_queue is None:
corr_uuid = uuid()
self._client_queue = '.'.join([command_name, 'client', corr_uuid])
def rpc(self,
command_name,
data=None,
server_routing_key=None):
"""Send a RPC request
:command_name: the command to execute (used as routing key)
:data: dict with data to be sent
:client_queue: Queue for this particular request
:server_routing_key: Server routing key. Will
default to <command>.server
"""
self.reply_received = False
payload = self._setup_payload(command_name, data)
logger.info("Preparing request {!r}".format(payload))
if server_routing_key is None:
if self._prefix is None:
server_routing_key = command_name
else:
server_routing_key = '.'.join([self._prefix, command_name])
self._prepare_client_queue(command_name)
logger.info("Set up client queue {!r} "
"to {!r}".format(self._client_queue,
server_routing_key))
message_correlation_id = uuid()
properties = {
'reply_to': self._client_queue,
'correlation_id': message_correlation_id
}
self.corr_id_server_queue[message_correlation_id] = server_routing_key
logger.info('STARTREQUEST:%s;CORRELATION_ID:%s' % (server_routing_key,
message_correlation_id))
result = None
try:
self.messages[message_correlation_id] = True
self._send_command(payload, server_routing_key, properties)
result = self.retrieve_messages()
except Exception:
logger.error("Sending message to consumer queue failed.",
exc_info=True)
raise
# Successful so store message correlation id for retrieval.
return result
def task(self,
command_name,
data=None,
server_routing_key=None):
"""Send a RPC request
:command_name: the command to execute (used as routing key)
:data: dict with data to be sent
:client_queue: Queue for this particular request
:server_routing_key: Server routing key. Will default
to <command>.server
"""
if data is None:
data = {}
self.reply_received = False
payload = self._setup_payload(command_name, data)
logger.debug("Preparing request {!r}".format(payload))
if server_routing_key is None:
server_routing_key = command_name
self._prepare_client_queue(command_name)
logger.info("Set up client queue {!r} "
"to {!r}".format(self._client_queue,
server_routing_key))
self._send_command(payload, server_routing_key, declare_queue=False)
def _send_command(self, payload, server_routing_key, properties=None,
declare_queue=True):
if properties is None:
properties = {}
self.reply = None
logger.debug("Using connection: {!r}".format(CONN_DICT))
logger.info("Declaring queue %s." % self._client_queue)
queue = Queue(self._client_queue,
channel=self._conn,
durable=self._exchange.durable,
exchange=self._exchange,
routing_key=self._client_queue)
if declare_queue:
queue.declare()
self._queue = queue
with producers[self._conn].acquire(block=True) as producer:
producer.publish(payload,
serializer='json',
exchange=self._exchange,
declare=[self._exchange],
routing_key=server_routing_key,
**properties)
logger.info("Published {!r} to exchange "
"{!r}".format(payload, self._exchange))
class Amqp(Thread):
def __init__(
self,
host="localhost",
port=5672,
userid="guest",
password="******",
virtual_host="canopsis",
exchange_name="canopsis",
logging_name="Amqp",
logging_level=INFO,
read_config_file=True,
auto_connect=True,
on_ready=None,
max_retries=5
):
super(Amqp, self).__init__()
self.logger = getLogger(logging_name)
self.host = host
self.port = port
self.userid = userid
self.password = password
self.virtual_host = virtual_host
self.exchange_name = exchange_name
self.logging_level = logging_level
if read_config_file:
self.read_config("amqp")
self.amqp_uri = "amqp://{0}:{1}@{2}:{3}/{4}".format(
self.userid, self.password, self.host, self.port,
self.virtual_host)
self.logger.setLevel(logging_level)
# Event sent try count before event drop in case of connection problem
if max_retries != 5:
self.logger.info('Custom retries value : {} {}'.format(
max_retries,
type(max_retries)
))
self.max_retries = max_retries
self.exchange_name_events = exchange_name + ".events"
self.exchange_name_alerts = exchange_name + ".alerts"
self.exchange_name_incidents = exchange_name + ".incidents"
self.chan = None
self.conn = None
self.connected = False
self.on_ready = on_ready
self.RUN = True
self.exchanges = {}
self.queues = {}
self.producers = None
self.paused = False
self.connection_errors = (
ConnectionError,
error,
IOError,
OSError)
# Create exchange
self.logger.debug("Create exchanges object")
for exchange_name in [
self.exchange_name, self.exchange_name_events,
self.exchange_name_alerts, self.exchange_name_incidents]:
self.logger.debug(' + {}'.format(exchange_name))
self.get_exchange(exchange_name)
if auto_connect:
self.connect()
self.logger.debug("Object canamqp initialized")
def run(self):
"""
main thread loop
- connects to rabbit
- calls drain_events
- disconnects on error
"""
self.logger.debug("Start thread ...")
reconnect = False
while self.RUN:
self.connect()
if self.connected:
self.init_queue(reconnect=reconnect)
self.logger.debug("Drain events ...")
while self.RUN:
try:
if not self.paused:
self.conn.drain_events(timeout=0.5)
else:
sleep(0.5)
except KombuSerializationError as exc:
self.logger.error(
"Kombu serialization error: invalid message "
"received: {}".format(exc))
except timeout:
pass
except self.connection_errors as err:
self.logger.error(
"Connection error ! ({})".format(err)
)
break
except KombuSerializationError as exc:
self.logger.error(
"Kombu serialization error: invalid message "
"received: {}".format(exc))
except Exception as err:
self.logger.error("Error: {} ({})".format(
err,
type(err)
))
print_exc(file=stdout)
break
self.disconnect()
if self.RUN:
self.logger.error(
'Connection lost, try to reconnect in few seconds ...'
)
reconnect = True
self.wait_connection(timeout=5)
self.logger.debug("End of thread ...")
def stop(self):
self.logger.debug("Stop thread ...")
self.RUN = False
def connect(self):
"""
Create the connection
Inits the producers
Init the only channel
Declares exchanges
"""
if not self.connected:
self.logger.info(
"Connect to AMQP Broker (%s:%s)" % (self.host, self.port))
self.conn = Connection(self.amqp_uri)
try:
self.logger.debug(" + Connect")
self.conn.connect()
self.logger.info("Connected to AMQP Broker.")
self.producers = pools.Producers(limit=10)
self.connected = True
except Exception as err:
self.conn.release()
self.logger.error("Impossible to connect ({})".format(err))
if self.connected:
self.logger.debug(" + Open channel")
try:
self.chan = self.conn.channel()
self.logger.info("Channel openned. Ready to send messages")
try:
# Declare exchange
self.logger.debug("Declare exchanges")
for exchange_name in self.exchanges:
self.logger.debug(" + {}".format(exchange_name))
self.exchanges[exchange_name](self.chan).declare()
except Exception as err:
self.logger.error(
"Impossible to declare exchange ({})".format(err))
except Exception as err:
self.logger.error(err)
else:
self.logger.debug("Already connected")
def get_exchange(self, name):
"""
Returns an exchange if stored in self.exchanges.
Otherwise, creates it and returns it
:param string name: name of the exchange to get/create
:rtype: Exchange|None
"""
if name:
try:
return self.exchanges[name]
except Exception:
if name == "amq.direct":
self.exchanges[name] = Exchange(
name, "direct", durable=True)
else:
self.exchanges[name] = Exchange(
name, "topic", durable=True, auto_delete=False)
return self.exchanges[name]
else:
return None
def init_queue(self, reconnect=False):
"""
Loads queue settings
Creates queues and attaches the same channel to each of them
Binds queues to exchange
Revives or creates consumers
calls consume
"""
if self.queues:
self.logger.debug("Init queues")
for queue_name in self.queues.keys():
self.logger.debug(" + {}".format(queue_name))
qsettings = self.queues[queue_name]
if not qsettings['queue']:
self.logger.debug(" + Create queue")
# copy list
routing_keys = list(qsettings['routing_keys'])
routing_key = None
if len(routing_keys):
routing_key = routing_keys[0]
routing_keys = routing_keys[1:]
exchange = self.get_exchange(qsettings['exchange_name'])
if qsettings['exchange_name'] == "amq.direct" \
and not routing_key:
routing_key = queue_name
self.logger.debug(
"exchange: '{}', exclusive: {},"
" auto_delete: {},no_ack: {}"
.format(
qsettings['exchange_name'],
qsettings['exclusive'],
qsettings['auto_delete'],
qsettings['no_ack']
)
)
qsettings['queue'] = Queue(
queue_name,
exchange=exchange,
routing_key=routing_key,
exclusive=qsettings['exclusive'],
auto_delete=qsettings['auto_delete'],
no_ack=qsettings['no_ack'],
channel=self.conn.channel())
qsettings['queue'].declare()
if len(routing_keys):
self.logger.debug(" + Bind on all routing keys")
for routing_key in routing_keys:
self.logger.debug(
" + routing_key: '{}'".format(routing_key)
)
try:
qsettings['queue'].bind_to(
exchange=exchange,
routing_key=routing_key
)
except Exception:
self.logger.error(
"You need upgrade your Kombu version ({})"
.format(__version__)
)
if qsettings['consumer'] and not reconnect:
qsettings['consumer'].revive(self.chan)
elif not qsettings['consumer'] or reconnect:
self.logger.debug(" + Create Consumer")
qsettings['consumer'] = self.conn.Consumer(
qsettings['queue'], callbacks=[qsettings['callback']])
self.logger.debug(" + Consume queue")
qsettings['consumer'].consume()
if self.on_ready:
self.on_ready()
else:
self.logger.info('Queue already inited')
def add_queue(
self,
queue_name,
routing_keys,
callback,
exchange_name=None,
no_ack=True,
exclusive=False,
auto_delete=True
):
"""
Initializes the queue configuration
maps the callback on the queue
"""
c_routing_keys = []
if not isinstance(routing_keys, list):
if isinstance(routing_keys, basestring):
c_routing_keys = [routing_keys]
else:
c_routing_keys = routing_keys
# aka if rk is nor a list nor a basetring, leave it empty
if not exchange_name:
exchange_name = self.exchange_name
self.queues[queue_name] = {
'queue': False,
'consumer': False,
'queue_name': queue_name,
'routing_keys': c_routing_keys,
'callback': callback,
'exchange_name': exchange_name,
'no_ack': no_ack,
'exclusive': exclusive,
'auto_delete': auto_delete
}
def publish(
self,
msg,
routing_key,
exchange_name=None,
serializer="json",
compression=None,
content_type=None,
content_encoding=None
):
"""
Tries to publish an event
In case of failure, tries to reconnect and retry until (self.max_retries)
:returns: operation success status
:rtype: bool
"""
self.logger.warning("Publishing from old.rabbitmq.Amqp is deprecated")
operation_success = False
retries = 0
while not operation_success and retries < self.max_retries:
retries += 1
if self.connected:
if not exchange_name:
exchange_name = self.exchange_name
with self.producers[self.conn].acquire(block=True) as producer:
try:
_msg = msg.copy()
Amqp._clean_msg_for_serialization(_msg)
exchange = self.get_exchange(
exchange_name.encode('utf-8'))
producer.publish(
_msg,
serializer=serializer,
compression=compression,
routing_key=routing_key,
exchange=exchange
)
self.logger.debug('publish {} in exchange {}'.format(
routing_key,
exchange_name
))
operation_success = True
except AmqpStructError:
self.logger.warning(
'Malformed message: routing key is too long. '
'Cancelling message')
return False
except Exception:
self.logger.error(
' + Impossible to send {}'
.format(traceback.format_exc())
)
self.disconnect()
self.connect()
self.init_queue(reconnect=False)
else:
self.logger.error('Not connected ... try reconnecting')
self.connect()
if not operation_success:
# Event and it's information are buffered until next send retry
self.logger.info('Retry count {}'.format(
retries
))
if not operation_success:
# Event and it's information are buffered until next send retry
self.logger.error('Too much retries for event {}, give up'.format(
routing_key
))
return operation_success
@staticmethod
def _clean_msg_for_serialization(msg):
from bson import objectid
for key in msg:
if isinstance(msg[key], objectid.ObjectId):
msg[key] = str(msg[key])
def cancel_queues(self):
if self.connected:
for queue_name in self.queues.keys():
if self.queues[queue_name]['consumer']:
self.logger.debug(
" + Cancel consumer on {}".format(queue_name)
)
try:
self.queues[queue_name]['consumer'].cancel()
except Exception:
pass
del(self.queues[queue_name]['consumer'])
self.queues[queue_name]['consumer'] = False
del(self.queues[queue_name]['queue'])
self.queues[queue_name]['queue'] = False
def disconnect(self):
if self.connected:
self.logger.info("Disconnect from AMQP Broker")
self.cancel_queues()
# Force producers closing to permit a clean reconnect after
# ... especially on timeout errors
self.producers[self.conn].force_close_all()
for exchange in self.exchanges:
del exchange
self.exchanges = {}
try:
pools.reset()
except Exception as err:
self.logger.error(
"Impossible to reset kombu pools: {} ({})".format(
err, type(err)))
try:
self.conn.release()
del self.conn
except Exception as err:
self.logger.error(
"Impossible to release connection: {} ({})".format(
err, type(err)))
self.connected = False
def wait_connection(self, timeout=5):
i = 0
while self.RUN and not self.connected and i < (timeout * 2):
sleep(0.5)
i += 1
def read_config(self, name):
filename = join(root_path, 'etc', '{0}.conf'.format(name))
import ConfigParser
self.config = ConfigParser.RawConfigParser()
try:
self.config.read(filename)
section = 'master'
self.host = self.config.get(section, "host")
self.port = self.config.getint(section, "port")
self.userid = self.config.get(section, "userid")
self.password = self.config.get(section, "password")
self.virtual_host = self.config.get(section, "virtual_host")
self.exchange_name = self.config.get(section, "exchange_name")
except Exception as err:
self.logger.error(
"Can't to load configurations ({}), use default ...".format(
err
))
def __del__(self):
self.stop()
class TriggerWatcher(ConsumerMixin):
sleep_interval = 4 # how long to sleep after processing each message
def __init__(self, create_handler, update_handler, delete_handler,
trigger_types=None, queue_suffix=None):
"""
:param create_handler: Function which is called on TriggerDB create event.
:type create_handler: ``callable``
:param update_handler: Function which is called on TriggerDB update event.
:type update_handler: ``callable``
:param delete_handler: Function which is called on TriggerDB delete event.
:type delete_handler: ``callable``
:param trigger_types: If provided, handler function will only be called
if the trigger in the message payload is included
in this list.
:type trigger_types: ``list``
"""
# TODO: Handle trigger type filtering using routing key
self._create_handler = create_handler
self._update_handler = update_handler
self._delete_handler = delete_handler
self._trigger_types = trigger_types
self._trigger_watch_q = self._get_queue(queue_suffix)
self.connection = None
self._load_thread = None
self._updates_thread = None
self._handlers = {
publishers.CREATE_RK: create_handler,
publishers.UPDATE_RK: update_handler,
publishers.DELETE_RK: delete_handler
}
def get_consumers(self, Consumer, channel):
return [Consumer(queues=[self._trigger_watch_q],
accept=['pickle'],
callbacks=[self.process_task])]
def process_task(self, body, message):
LOG.debug('process_task')
LOG.debug(' body: %s', body)
LOG.debug(' message.properties: %s', message.properties)
LOG.debug(' message.delivery_info: %s', message.delivery_info)
routing_key = message.delivery_info.get('routing_key', '')
handler = self._handlers.get(routing_key, None)
try:
if not handler:
LOG.debug('Skipping message %s as no handler was found.', message)
return
trigger_type = getattr(body, 'type', None)
if self._trigger_types and trigger_type not in self._trigger_types:
LOG.debug('Skipping message %s since\'t trigger_type doesn\'t match (type=%s)',
message, trigger_type)
return
try:
handler(body)
except Exception as e:
LOG.exception('Handling failed. Message body: %s. Exception: %s',
body, e.message)
finally:
message.ack()
eventlet.sleep(self.sleep_interval)
def start(self):
try:
self.connection = Connection(cfg.CONF.messaging.url)
self._updates_thread = eventlet.spawn(self.run)
self._load_thread = eventlet.spawn(self._load_triggers_from_db)
except:
LOG.exception('Failed to start watcher.')
self.connection.release()
def stop(self):
try:
self._updates_thread = eventlet.kill(self._updates_thread)
self._load_thread = eventlet.kill(self._load_thread)
finally:
self.connection.release()
# Note: We sleep after we consume a message so we give a chance to other
# green threads to run. If we don't do that, ConsumerMixin will block on
# waiting for a message on the queue.
def on_consume_end(self, connection, channel):
super(TriggerWatcher, self).on_consume_end(connection=connection,
channel=channel)
eventlet.sleep(seconds=self.sleep_interval)
def on_iteration(self):
super(TriggerWatcher, self).on_iteration()
eventlet.sleep(seconds=self.sleep_interval)
def _load_triggers_from_db(self):
for trigger_type in self._trigger_types:
for trigger in Trigger.query(type=trigger_type):
LOG.debug('Found existing trigger: %s in db.' % trigger)
self._handlers[publishers.CREATE_RK](trigger)
@staticmethod
def _get_queue(queue_suffix):
if not queue_suffix:
# pick last 10 digits of uuid. Arbitrary but unique enough for the TriggerWatcher.
u_hex = uuid.uuid4().hex
queue_suffix = uuid.uuid4().hex[len(u_hex) - 10:]
queue_name = 'st2.trigger.watch.%s' % queue_suffix
return reactor.get_trigger_cud_queue(queue_name, routing_key='#')
class camqp(threading.Thread):
def __init__(self, host="localhost", port=5672, userid="guest", password="******", virtual_host="canopsis", exchange_name="canopsis", logging_name="camqp", logging_level=logging.INFO, read_config_file=True, auto_connect=True, on_ready=None):
threading.Thread.__init__(self)
self.logger = logging.getLogger(logging_name)
self.host=host
self.port=port
self.userid=userid
self.password=password
self.virtual_host=virtual_host
self.exchange_name=exchange_name
self.logging_level = logging_level
if (read_config_file):
self.read_config("amqp")
self.amqp_uri = "amqp://%s:%s@%s:%s/%s" % (self.userid, self.password, self.host, self.port, self.virtual_host)
self.logger.setLevel(logging_level)
self.exchange_name_events=exchange_name+".events"
self.exchange_name_alerts=exchange_name+".alerts"
self.exchange_name_incidents=exchange_name+".incidents"
self.chan = None
self.conn = None
self.connected = False
self.on_ready = on_ready
self.RUN = True
self.exchanges = {}
self.queues = {}
self.paused = False
self.connection_errors = ( ConnectionError,
socket.error,
IOError,
OSError,)
#AttributeError)
## create exchange
self.logger.debug("Create exchanges object")
for exchange_name in [self.exchange_name, self.exchange_name_events, self.exchange_name_alerts, self.exchange_name_incidents]:
self.logger.debug(" + %s" % exchange_name)
self.get_exchange(exchange_name)
if auto_connect:
self.connect()
self.logger.debug("Object canamqp initialized")
def run(self):
self.logger.debug("Start thread ...")
reconnect = False
while self.RUN:
self.connect()
#self.wait_connection()
if self.connected:
self.init_queue(reconnect=reconnect)
self.logger.debug("Drain events ...")
while self.RUN:
try:
if not self.paused:
self.conn.drain_events(timeout=0.5)
else:
time.sleep(0.5)
except socket.timeout:
pass
except self.connection_errors as err:
self.logger.error("Connection error ! (%s)" % err)
break
except Exception as err:
self.logger.error("Unknown error: %s (%s)" % (err, type(err)))
traceback.print_exc(file=sys.stdout)
break
self.disconnect()
if self.RUN:
self.logger.error("Connection loss, try to reconnect in few seconds ...")
reconnect = True
self.wait_connection(timeout=5)
self.logger.debug("End of thread ...")
def stop(self):
self.logger.debug("Stop thread ...")
self.RUN = False
def connect(self):
if not self.connected:
self.logger.info("Connect to AMQP Broker (%s:%s)" % (self.host, self.port))
self.conn = Connection(self.amqp_uri)
try:
self.logger.debug(" + Connect")
self.conn.connect()
self.logger.info("Connected to AMQP Broker.")
self.producers = kombu.pools.Producers(limit=10)
self.connected = True
except Exception as err:
self.conn.release()
self.logger.error("Impossible to connect (%s)" % err)
if self.connected:
self.logger.debug(" + Open channel")
try:
self.chan = self.conn.channel()
self.logger.debug("Channel openned. Ready to send messages")
try:
## declare exchange
self.logger.debug("Declare exchanges")
for exchange_name in self.exchanges:
self.logger.debug(" + %s" % exchange_name)
self.exchanges[exchange_name](self.chan).declare()
except Exception as err:
self.logger.error("Impossible to declare exchange (%s)" % err)
except Exception as err:
self.logger.error(err)
else:
self.logger.debug("Allready connected")
def get_exchange(self, name):
if name:
try:
return self.exchanges[name]
except:
if name == "amq.direct":
self.exchanges[name] = Exchange(name, "direct", durable=True)
else:
self.exchanges[name] = Exchange(name , "topic", durable=True, auto_delete=False)
return self.exchanges[name]
else:
return None
def init_queue(self, reconnect=False):
if self.queues:
self.logger.debug("Init queues")
for queue_name in self.queues.keys():
self.logger.debug(" + %s" % queue_name)
qsettings = self.queues[queue_name]
if not qsettings['queue']:
self.logger.debug(" + Create queue")
# copy list
routing_keys = list(qsettings['routing_keys'])
routing_key = None
if len(routing_keys):
routing_key = routing_keys[0]
routing_keys = routing_keys[1:]
exchange = self.get_exchange(qsettings['exchange_name'])
if (qsettings['exchange_name'] == "amq.direct" and not routing_key):
routing_key = queue_name
#self.logger.debug(" + exchange: '%s', routing_key: '%s', exclusive: %s, auto_delete: %s, no_ack: %s" % (qsettings['exchange_name'], routing_key, qsettings['exclusive'], qsettings['auto_delete'], qsettings['no_ack']))
self.logger.debug(" + exchange: '%s', exclusive: %s, auto_delete: %s, no_ack: %s" % (qsettings['exchange_name'], qsettings['exclusive'], qsettings['auto_delete'], qsettings['no_ack']))
qsettings['queue'] = Queue(queue_name,
exchange = exchange,
routing_key = routing_key,
exclusive = qsettings['exclusive'],
auto_delete = qsettings['auto_delete'],
no_ack = qsettings['no_ack'],
channel=self.conn.channel())
qsettings['queue'].declare()
if len(routing_keys):
self.logger.debug(" + Bind on all routing keys")
for routing_key in routing_keys:
self.logger.debug(" + routing_key: '%s'" % routing_key)
try:
qsettings['queue'].bind_to(exchange=exchange, routing_key=routing_key)
except:
self.logger.error("You need upgrade your Kombu version (%s)" % kombu.__version__)
if not qsettings['consumer']:
self.logger.debug(" + Create Consumer")
qsettings['consumer'] = self.conn.Consumer(qsettings['queue'], callbacks=[ qsettings['callback'] ])
self.logger.debug(" + Consume queue")
qsettings['consumer'].consume()
if self.on_ready:
self.on_ready()
def add_queue(self, queue_name, routing_keys, callback, exchange_name=None, no_ack=True, exclusive=False, auto_delete=True):
#if exchange_name == "amq.direct":
# routing_keys = queue_name
c_routing_keys = []
if not isinstance(routing_keys, list):
if isinstance(routing_keys, str):
c_routing_keys = [ routing_keys ]
else:
c_routing_keys = routing_keys
if not exchange_name:
exchange_name = self.exchange_name
self.queues[queue_name]={ 'queue': False,
'consumer': False,
'queue_name': queue_name,
'routing_keys': c_routing_keys,
'callback': callback,
'exchange_name': exchange_name,
'no_ack': no_ack,
'exclusive': exclusive,
'auto_delete': auto_delete
}
def publish(self, msg, routing_key, exchange_name=None, serializer="json", compression=None, content_type=None, content_encoding=None):
self.wait_connection()
if self.connected:
if not exchange_name:
exchange_name = self.exchange_name
self.logger.debug("Send message to %s in %s" % (routing_key, exchange_name))
with self.producers[self.conn].acquire(block=True) as producer:
try:
_msg = msg.copy()
camqp._clean_msg_for_serialization(_msg)
producer.publish(_msg, serializer=serializer, compression=compression, routing_key=routing_key, exchange=self.get_exchange(exchange_name))
self.logger.debug(" + Sended")
except Exception, err:
self.logger.error(" + Impossible to send (%s)" % err)
else:
def do_consume(user_qs):
print("about to listen no queues [%s]" %
", ".join(list(map(lambda x: x, user_qs))))
conn = Connection(amqp_hosts, failover_strategy='round-robin')
# try to get a connection no matter what
while True:
try:
conn.ensure_connection(errback=on_ens_conn_err_cb)
conn.connect()
except Exception as e:
print("connection error failed on exception [%s]" % repr(e))
conn.release()
continue
if conn.connected:
break
else:
print("connection failed in some way, retry")
chan = conn.channel()
global bound_cons_Q
cons_Q = Queue(common.uuid(), queue_arguments=q_expires)
bound_cons_Q = cons_Q(chan)
bound_cons_Q.declare()
# first bind to some control route
bound_cons_Q.bind_to(priTopicXchg, routing_key='manage.#')
for i in user_qs:
if '*' in i or '#' in i:
# create the wildcard route_key bind
bound_cons_Q.bind_to(priTopicXchg, routing_key=i)
else:
for j in allQs:
if i == j.as_dict()['name']:
bound_cons_Q.bind_to(priTopicXchg, routing_key=j.as_dict()['routing_key'])
cons = Consumer(
chan,
accept=['json'],
queues=bound_cons_Q,
callbacks=[on_msg_cb_1, on_msg_cb_2]
)
print("queue set to [%s]" % bound_cons_Q.as_dict(recurse=True))
cons.consume()
while True:
try:
conn.drain_events()
except conn.connection_errors + conn.channel_errors as e:
print("connection [%s] went down (error[%s]), trying to "
"connect to the next one" % (conn.info(), repr(e)))
conn.close()
conn.release()
conn.ensure_connection(errback=on_ens_conn_err_cb)
conn.connect()
chan = conn.channel()
cons_Q.bind(chan)
cons = Consumer(
chan,
accept=['json'],
queues=bound_cons_Q,
callbacks=[on_msg_cb_1, on_msg_cb_2]
)
cons.consume()
class RabbitMQSource(StoqSourcePlugin):
def __init__(self):
super().__init__()
def activate(self, stoq):
self.stoq = stoq
super().activate()
self.amqp_publish_conn = None
if self.use_ssl:
try:
self.ssl_config = {}
self.ssl_config['ca_certs'] = self.ca_certs
self.ssl_config['keyfile'] = self.keyfile
self.ssl_config['certfile'] = self.certfile
if self.cert_reqs:
self.ssl_config['cert_reqs'] = ssl.CERT_REQUIRED
except AttributeError as err:
self.log.critical("SSL option missing, but required: {}".format(err))
self.ssl_config = False
else:
self.ssl_config = False
def ingest(self):
"""
Monitor AMQP for messages
"""
# Define our RabbitMQ route
routing_key = self.stoq.worker.name
# If this is an error message, let's make sure our queue
# has "-errors" affixed to it
if self.stoq.worker.error_queue is True:
routing_key = routing_key + "-errors".strip()
exchange = Exchange(self.exchange_name, type=self.exchange_type)
queue_arguments = {'x-max-priority': 10}
queue = Queue(routing_key, exchange, routing_key=routing_key,
queue_arguments=queue_arguments)
self.log.info("Monitoring {} queue for messages...".format(routing_key))
# Setup our broker connection with RabbitMQ
with Connection(hostname=self.host,
port=self.port,
userid=self.user,
password=self.password,
virtual_host=self.virtual_host,
ssl=self.ssl_config) as conn:
conn.connect()
consumer = Consumer(conn, queue, callbacks=[self.queue_callback])
consumer.qos(prefetch_count=int(self.prefetch))
consumer.consume()
while True:
try:
conn.drain_events()
except Exception as err:
self.log.critical("Unable to process queue: {}".format(err), exc_info=True)
conn.release()
break
def queue_callback(self, amqp_message_data, amqp_message_handler):
try:
# Setup the amqp message for parsing
kwargs = self.stoq.loads(amqp_message_data)
# Send the message to the worker
self.stoq.worker.multiprocess_put(**kwargs)
except Exception as err:
self.log.error("Unable to process message: {}, msg: {}".format(err, kwargs))
# Something went wrong, let's publish to the error queue. and
# append the error message.
kwargs['err'] = str(err)
self.publish_connect()
self.publish(kwargs, self.stoq.worker.name, err=True)
self.publish_release()
# Acknowledge our message in amqp so we can get the next one.
amqp_message_handler.ack()
def publish_connect(self):
"""
Connect to AMQP to publish a message
:returns: AMQP connection object for publishing
:rtype: kombu.Connection object
"""
self.amqp_exchange = Exchange(self.exchange_name,
type=self.exchange_type)
self.amqp_publish_conn = Connection(hostname=self.host,
port=self.port,
userid=self.user,
password=self.password,
virtual_host=self.virtual_host,
ssl=self.ssl_config)
self.amqp_publish_conn.connect()
return self.amqp_publish_conn
def publish_release(self):
"""
Release AMQP connection used for publishing
"""
return self.amqp_publish_conn.release()
def publish(self, msg, routing_key, err=False, **kwargs):
"""
Publish a message to AMQP
:param dict msg: Message to be published
:param str routing_key: Route to be used, should be name of worker
:param int priority: Priority of the message, between 0-10
:param bool err: Define whether we should process error queue
"""
# Make sure we have a valid connection to RabbitMQ
if not self.amqp_publish_conn:
self.publish_connect()
priority = kwargs.get('priority', 0)
# If this is an error message, let's make sure our queue
# has "-errors" affixed to it
if err:
routing_key = routing_key + "-errors".strip()
# Define the queue so we can ensure it is declared before
# publishing to it
queue_arguments = {'x-max-priority': 10}
queue = Queue(routing_key, self.amqp_exchange, routing_key=routing_key,
queue_arguments=queue_arguments)
with producers[self.amqp_publish_conn].acquire(block=True) as producer:
producer.publish(self.stoq.dumps(msg),
exchange=self.amqp_exchange,
declare=[self.amqp_exchange, queue],
routing_key=routing_key,
priority=priority)
def __errback(self, exc, interval):
"""
Error handling for AMQP publishing
"""
self.log.warn("Unable to publish message: {}. Retry in {}s.".format(exc, interval))
class PivWorker(ConsumerMixin):
def __init__(self, contributor):
if contributor.connector_type != ConnectorType.piv.value:
raise ValueError(
"Contributor '{0}': PivWorker requires type {1}".format(contributor.id, ConnectorType.piv.value)
)
if not contributor.is_active:
raise ValueError(
"Contributor '{0}': PivWorker requires an activated contributor.".format(contributor.id)
)
if not contributor.broker_url:
raise ValueError("Missing 'broker_url' configuration for contributor '{0}'".format(contributor.id))
if not contributor.exchange_name:
raise ValueError(
"Missing 'exchange_name' configuration for contributor '{0}'".format(contributor.id)
)
if not contributor.queue_name:
raise ValueError("Missing 'queue_name' configuration for contributor '{0}'".format(contributor.id))
self.last_config_checked_time = datetime.now()
self.broker_url = deepcopy(contributor.broker_url)
self.builder = KirinModelBuilder(contributor)
def __enter__(self):
self.connection = Connection(self.builder.contributor.broker_url)
self.exchange = self._get_exchange(self.builder.contributor.exchange_name)
self.queue = self._get_queue(self.exchange, self.builder.contributor.queue_name)
return self
def __exit__(self, type, value, traceback):
self.connection.release()
def _get_exchange(self, exchange_name):
return Exchange(exchange_name, "fanout", durable=True, no_declare=True)
def _get_queue(self, exchange, queue_name):
return Queue(queue_name, exchange, durable=True, auto_delete=False)
def get_consumers(self, Consumer, channel):
return [
Consumer(
queues=[self.queue],
accept=["application/json"],
prefetch_count=1,
callbacks=[self.process_message],
)
]
def process_message(self, body, message):
wrap_build(self.builder, body)
# TODO: We might want to not acknowledge the message in case of error in
# the processing.
message.ack()
def on_iteration(self):
if datetime.now() - self.last_config_checked_time < CONF_RELOAD_INTERVAL:
return
else:
# SQLAlchemy is not querying the DB for read (uses cache instead),
# unless we specifically tell that the data is expired.
db.session.expire(self.builder.contributor, ["broker_url", "exchange_name", "queue_name"])
self.last_config_checked_time = datetime.now()
contributor = get_piv_contributor(self.builder.contributor.id)
if not contributor:
logger.info(
"contributor '{0}' doesn't exist anymore, let the worker die".format(self.builder.contributor.id)
)
self.should_stop = True
return
if contributor.broker_url != self.broker_url:
logger.info("broker URL for contributor '{0}' changed, let the worker die".format(contributor.id))
self.should_stop = True
return
if contributor.exchange_name != self.exchange.name:
logger.info(
"exchange name for contributor '{0}' changed, worker updated".format(contributor.exchange_name)
)
self.exchange = self._get_exchange(contributor.exchange_name)
self.queue = self._get_queue(self.exchange, contributor.queue_name)
if contributor.queue_name != self.queue.name:
logger.info(
"queue name for contributor '{0}' changed, worker updated".format(contributor.queue_name)
)
self.queue = self._get_queue(self.exchange, contributor.queue_name)
self.builder = KirinModelBuilder(contributor)
class GenericConsumer(object):
"""Generic publisher class that specific consumers inherit from.
FIXME: Mandatory properties, like "topic", should not be set from generic
functions like configure() but should instead be explicitly required
somewhere, e.g. the constructor.
"""
def __init__(self, config, exchange=None, connect=True, heartbeat=False,
**kwargs):
self.config = config
self.exchange = exchange
self.connection = None
self.durable = False
self.applabel = ''
self.heartbeat = heartbeat
for x in ['applabel', 'topic', 'callback', 'durable']:
if x in kwargs:
setattr(self, x, kwargs[x])
del kwargs[x]
# Only used if there is no applabel.
self.queue_gen_name = None
if connect:
self.connect()
@property
def queue_name(self):
# This is a property instead of being set in the constructor since
# applabel can be set later via configure().
queue_name = 'queue/%s/' % self.config.user
if self.applabel:
queue_name += self.applabel
else:
if not self.queue_gen_name:
self.queue_gen_name = str(uuid.uuid1())
queue_name += self.queue_gen_name
return queue_name
def configure(self, **kwargs):
"""Sets variables."""
for x in kwargs:
setattr(self, x, kwargs[x])
def connect(self):
if not self.connection:
self.connection = Connection(hostname=self.config.host,
port=self.config.port,
userid=self.config.user,
password=self.config.password,
virtual_host=self.config.vhost,
ssl=self.config.ssl)
def disconnect(self):
if self.connection:
self.connection.release()
self.connection = None
def purge_existing_messages(self):
"""Purge messages that are already in the queue on the broker.
TODO: I think this is only supported by the amqp backend.
"""
if self.durable and not self.applabel:
raise InvalidAppLabel('Durable consumers must have an applabel')
if not self.connection:
self.connect()
queue = self._create_queue()
try:
queue(self.connection).purge()
except ChannelError, e:
if e.message == 404:
pass
raise
from kombu import Connection
connection = Connection('amqp://*****:*****@localhost:5672//')
connection.connect()
#check whether connected boolean
connection.connected
connection.close()
#better practice, releases resource if in a connection pool else close
connection.release()
class Amqp(Thread):
def __init__(
self,
host="localhost",
port=5672,
userid="guest",
password="******",
virtual_host="canopsis",
exchange_name="canopsis",
logging_name="Amqp",
logging_level=INFO,
read_config_file=True,
auto_connect=True,
on_ready=None,
max_retries=5,
):
super(Amqp, self).__init__()
self.logger = getLogger(logging_name)
self.host = host
self.port = port
self.userid = userid
self.password = password
self.virtual_host = virtual_host
self.exchange_name = exchange_name
self.logging_level = logging_level
if read_config_file:
self.read_config("amqp")
self.amqp_uri = "amqp://{0}:{1}@{2}:{3}/{4}".format(
self.userid, self.password, self.host, self.port, self.virtual_host
)
self.logger.setLevel(logging_level)
# Event sent try count before event drop in case of connection problem
if max_retries != 5:
self.logger.info(u"Custom retries value : {} {}".format(max_retries, type(max_retries)))
self.max_retries = max_retries
self.exchange_name_events = exchange_name + ".events"
self.exchange_name_alerts = exchange_name + ".alerts"
self.exchange_name_incidents = exchange_name + ".incidents"
self.chan = None
self.conn = None
self.connected = False
self.on_ready = on_ready
self.RUN = True
self.exchanges = {}
self.queues = {}
self.paused = False
self.connection_errors = (ConnectionError, error, IOError, OSError)
# Create exchange
self.logger.debug("Create exchanges object")
for exchange_name in [
self.exchange_name,
self.exchange_name_events,
self.exchange_name_alerts,
self.exchange_name_incidents,
]:
self.logger.debug(u" + {0}".format(exchange_name))
self.get_exchange(exchange_name)
if auto_connect:
self.connect()
self.logger.debug("Object canamqp initialized")
def run(self):
self.logger.debug("Start thread ...")
reconnect = False
while self.RUN:
self.connect()
if self.connected:
self.init_queue(reconnect=reconnect)
self.logger.debug("Drain events ...")
while self.RUN:
try:
if not self.paused:
self.conn.drain_events(timeout=0.5)
else:
sleep(0.5)
except timeout:
pass
except self.connection_errors as err:
self.logger.error(u"Connection error ! ({})".format(err))
break
except Exception as err:
self.logger.error(u"Error: {} ({})".format(err, type(err)))
print_exc(file=stdout)
break
self.disconnect()
if self.RUN:
self.logger.error("Connection lost, try to reconnect in few seconds ...")
reconnect = True
self.wait_connection(timeout=5)
self.logger.debug("End of thread ...")
def stop(self):
self.logger.debug("Stop thread ...")
self.RUN = False
def connect(self):
if not self.connected:
self.logger.info("Connect to AMQP Broker (%s:%s)" % (self.host, self.port))
self.conn = Connection(self.amqp_uri)
try:
self.logger.debug(" + Connect")
self.conn.connect()
self.logger.info("Connected to AMQP Broker.")
self.producers = pools.Producers(limit=10)
self.connected = True
except Exception as err:
self.conn.release()
self.logger.error(u"Impossible to connect ({})".format(err))
if self.connected:
self.logger.debug(" + Open channel")
try:
self.chan = self.conn.channel()
self.logger.debug("Channel openned. Ready to send messages")
try:
# Declare exchange
self.logger.debug("Declare exchanges")
for exchange_name in self.exchanges:
self.logger.debug(u" + {}".format(exchange_name))
self.exchanges[exchange_name](self.chan).declare()
except Exception as err:
self.logger.error(u"Impossible to declare exchange ({})".format(err))
except Exception as err:
self.logger.error(err)
else:
self.logger.debug("Already connected")
def get_exchange(self, name):
if name:
try:
return self.exchanges[name]
except:
if name == "amq.direct":
self.exchanges[name] = Exchange(name, "direct", durable=True)
else:
self.exchanges[name] = Exchange(name, "topic", durable=True, auto_delete=False)
return self.exchanges[name]
else:
return None
def init_queue(self, reconnect=False):
if self.queues:
self.logger.debug("Init queues")
for queue_name in self.queues.keys():
self.logger.debug(u" + {}".format(queue_name))
qsettings = self.queues[queue_name]
if not qsettings["queue"]:
self.logger.debug(" + Create queue")
# copy list
routing_keys = list(qsettings["routing_keys"])
routing_key = None
if len(routing_keys):
routing_key = routing_keys[0]
routing_keys = routing_keys[1:]
exchange = self.get_exchange(qsettings["exchange_name"])
if qsettings["exchange_name"] == "amq.direct" and not routing_key:
routing_key = queue_name
self.logger.debug(
(u"exchange: '{}', exclusive: {}," + " auto_delete: {},no_ack: {}").format(
qsettings["exchange_name"],
qsettings["exclusive"],
qsettings["auto_delete"],
qsettings["no_ack"],
)
)
qsettings["queue"] = Queue(
queue_name,
exchange=exchange,
routing_key=routing_key,
exclusive=qsettings["exclusive"],
auto_delete=qsettings["auto_delete"],
no_ack=qsettings["no_ack"],
channel=self.conn.channel(),
)
qsettings["queue"].declare()
if len(routing_keys):
self.logger.debug(" + Bind on all routing keys")
for routing_key in routing_keys:
self.logger.debug(" + routing_key: '{}'".format(routing_key))
try:
qsettings["queue"].bind_to(exchange=exchange, routing_key=routing_key)
except:
self.logger.error(u"You need upgrade your Kombu version ({})".format(__version__))
if not qsettings["consumer"] or reconnect:
self.logger.debug(" + Create Consumer")
qsettings["consumer"] = self.conn.Consumer(qsettings["queue"], callbacks=[qsettings["callback"]])
self.logger.debug(" + Consume queue")
qsettings["consumer"].consume()
if self.on_ready:
self.on_ready()
def add_queue(
self, queue_name, routing_keys, callback, exchange_name=None, no_ack=True, exclusive=False, auto_delete=True
):
c_routing_keys = []
if not isinstance(routing_keys, list):
if isinstance(routing_keys, basestring):
c_routing_keys = [routing_keys]
else:
c_routing_keys = routing_keys
if not exchange_name:
exchange_name = self.exchange_name
self.queues[queue_name] = {
"queue": False,
"consumer": False,
"queue_name": queue_name,
"routing_keys": c_routing_keys,
"callback": callback,
"exchange_name": exchange_name,
"no_ack": no_ack,
"exclusive": exclusive,
"auto_delete": auto_delete,
}
def publish(
self,
msg,
routing_key,
exchange_name=None,
serializer="json",
compression=None,
content_type=None,
content_encoding=None,
):
operation_success = False
retries = 0
while not operation_success and retries < self.max_retries:
retries += 1
if self.connected:
if not exchange_name:
exchange_name = self.exchange_name
with self.producers[self.conn].acquire(block=True) as producer:
try:
_msg = msg.copy()
Amqp._clean_msg_for_serialization(_msg)
producer.publish(
_msg,
serializer=serializer,
compression=compression,
routing_key=routing_key,
exchange=self.get_exchange(exchange_name),
)
self.logger.debug("publish {} in exchange {}".format(routing_key, exchange_name))
operation_success = True
except Exception as e:
self.logger.error(u" + Impossible to send {}".format(e))
self.disconnect()
self.connect()
else:
self.logger.error("Not connected ... try reconnecting")
self.connect()
if not operation_success:
# Event and it's information are buffered until next send retry
self.logger.info(u"Retry count {}".format(retries))
if not operation_success:
# Event and it's information are buffered until next send retry
self.logger.error(u"Too much retries for event {}, give up".format(routing_key))
@staticmethod
def _clean_msg_for_serialization(msg):
from bson import objectid
for key in msg:
if isinstance(msg[key], objectid.ObjectId):
msg[key] = str(msg[key])
def cancel_queues(self):
if self.connected:
for queue_name in self.queues.keys():
if self.queues[queue_name]["consumer"]:
self.logger.debug(u" + Cancel consumer on {}".format(queue_name))
try:
self.queues[queue_name]["consumer"].cancel()
except:
pass
del (self.queues[queue_name]["consumer"])
self.queues[queue_name]["consumer"] = False
del (self.queues[queue_name]["queue"])
self.queues[queue_name]["queue"] = False
def disconnect(self):
if self.connected:
self.logger.info("Disconnect from AMQP Broker")
self.cancel_queues()
for exchange in self.exchanges:
del exchange
self.exchanges = {}
try:
pools.reset()
except Exception as err:
self.logger.error(u"Impossible to reset kombu pools: {} ({})".format(err, type(err)))
try:
self.conn.release()
del self.conn
except Exception as err:
self.logger.error(u"Impossible to release connection: {} ({})".format(err, type(err)))
self.connected = False
def wait_connection(self, timeout=5):
i = 0
while self.RUN and not self.connected and i < (timeout * 2):
try:
sleep(0.5)
except:
pass
i += 1
def read_config(self, name):
filename = join(sys_prefix, "etc", u"{0}.conf".format(name))
import ConfigParser
self.config = ConfigParser.RawConfigParser()
try:
self.config.read(filename)
section = "master"
self.host = self.config.get(section, "host")
self.port = self.config.getint(section, "port")
self.userid = self.config.get(section, "userid")
self.password = self.config.get(section, "password")
self.virtual_host = self.config.get(section, "virtual_host")
self.exchange_name = self.config.get(section, "exchange_name")
except Exception as err:
self.logger.error(u"Can't to load configurations ({}), use default ...".format(err))
def __del__(self):
self.stop()
class AutophonePulseMonitor(object):
"""AutophonePulseMonitor provides the means to be notified when
Android builds are available for testing and when users have initiated
retriggers and cancels via the Treeherder UI. Builds can be selected using
repository names, Android platform names or build types.
AutophonePulseMonitor detects new builds by listening to
un-normalized buildbot initiated pulse messages rather than the
normalized messages in order to obtain the check-in comment for a
build. The comment is used to determine if a try build has
requested Autophone testing.
:param hostname: Hostname of Pulse. Defaults to the production pulse
server pulse.mozilla.org.
:param userid: Pulse User id
:param password: Pulse Password
:param virtual_host: AMQP virtual host, defaults to '/'.
:param durable_queues: If True, will create durable queues in
Pulse for the build and job action messages. Defaults to
False. In production, durable_queues should be set to True to
avoid losing messages if the connection is broken or the
application crashes.
:param build_exchange_name: Name of build exchange. Defaults to
'exchange/build/'.
:param build_queue_name: Build queue name suffix. Defaults to
'builds'. The pulse build queue will be created with a name
of the form 'queue/<userid>/<build_queue_name>'.
:param jobaction_exchange_name: Name of job action exchange.
Defaults to 'exchange/treeherder/v1/job-actions'. Use
'exchange/treeherder-stage/v1/job-actions' to listen to job
action messages for Treeherder staging.
:param jobaction_queue_name: Job action queue name suffix. Defaults to
'jobactions'. The pulse jobaction queue will be created with a name
of the form 'queue/<userid>/<jobaction_queue_name>'.
:param build_callback: Required callback function which takes a
single `build_data` object as argument containing information
on matched builds. `build_callback` is always called on a new
thread. `build_data` is an object which is guaranteed to
contain the following keys:
'appName': Will always be 'Fennec'
'branch': The repository name of the build, e.g. 'mozilla-central'.
'comments': Check-in comment.
'packageUrl': The url to the apk package for the build.
'platform': The platform name of the build, e.g. 'android-api-11'
`build_data` may also contain the following keys:
'buildid': Build id in CCYYMMDDHHMMSS format.
'robocopApkUrl': Url to robocop apk for the build.
'symbolsUrl': Url to the symbols zip file for the build.
'testsUrl': Url to the tests zip file for the build.
'who': Check-in Commiter.
:param jobaction_callback: Required callback function which takes a
single `jobaction_data` object as argument containing information
on matched actions. `jobaction_callback` is always called on a new
thread. `jobaction_data` is an object which is contains the following keys:
'action': 'cancel' or 'retrigger',
'project': repository name,
'job_id': treeherder job_id,
'job_guid': treeherder job_guid,
'build_type': 'opt' or 'debug',
'platform': the detected platform,
'build_url': build url,
'machine_name': name of machine ,
'job_group_name': treeherder job group name,
'job_group_symbol': treeherder job group symbol,
'job_type_name': treeherder job type name,
'job_type_symbol': treeherder job type symbol,
'result': test result result',
:param treeherder_url: Optional Treeherder server url if Treeherder
job action pulse messages are to be processed. Defaults to None.
:param trees: Required list of repository names to be matched.
:param platforms: Required list of platforms to be
matched. Currently, the possible values are 'android',
'android-api-9', 'android-api-10', 'android-api-11', and
'android-x86'.
:param buildtypes: Required list of build types to
process. Possible values are 'opt', 'debug'
:param timeout: Timeout in seconds for the kombu connection
drain_events. Defaults to 5 seconds.
:param shared_lock: Required lock used to control concurrent
access. Used to prevent socket based deadlocks.
:param verbose: If True, will log build and job action messages.
Defaults to False.
Usage:
::
import threading
import time
from optparse import OptionParser
parser = OptionParser()
def build_callback(build_data):
logger = logging.getLogger()
logger.debug('PULSE BUILD FOUND %s' % build_data)
def jobaction_callback(job_action):
logger = logging.getLogger()
if job_action['job_group_name'] != 'Autophone':
return
logger.debug('JOB ACTION FOUND %s' % json.dumps(
job_action, sort_keys=True, indent=4))
logging.basicConfig()
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
parser.add_option('--pulse-user', action='store', type='string',
dest='pulse_user', default='',
help='user id for connecting to PulseGuardian')
parser.add_option('--pulse-password', action='store', type='string',
dest='pulse_password', default='',
help='password for connecting to PulseGuardian')
(options, args) = parser.parse_args()
shared_lock = threading.Lock()
monitor = AutophonePulseMonitor(
userid=options.pulse_user,
password=options.pulse_password,
jobaction_exchange_name='exchange/treeherder-stage/v1/job-actions',
build_callback=build_callback,
jobaction_callback=jobaction_callback,
trees=['try', 'mozilla-inbound'],
platforms=['android-api-9', 'android-api-11'],
buildtypes=['opt'],
shared_lock=shared_lock)
monitor.start()
time.sleep(3600)
"""
def __init__(self,
hostname='pulse.mozilla.org',
userid=None,
password=None,
virtual_host='/',
durable_queues=False,
build_exchange_name='exchange/build/',
build_queue_name='builds',
jobaction_exchange_name='exchange/treeherder/v1/job-actions',
jobaction_queue_name='jobactions',
build_callback=None,
jobaction_callback=None,
treeherder_url=None,
trees=[],
platforms=[],
buildtypes=[],
timeout=5,
shared_lock=None,
verbose=False):
assert userid, "userid is required."
assert password, "password is required."
assert build_callback, "build_callback is required."
assert trees, "trees is required."
assert platforms, "platforms is required."
assert buildtypes, "buildtypes is required."
assert shared_lock, "shared_lock is required."
self.treeherder_url = treeherder_url
self.build_callback = build_callback
self.jobaction_callback = jobaction_callback
self.trees = list(trees)
self.platforms = list(platforms)
# Sort the platforms in descending order of length, so we do
# not make a match on a substring of the platform prematurely.
self.platforms.sort(cmp=lambda x,y: (len(y) - len(x)))
self.buildtypes = list(buildtypes)
self.timeout = timeout
self.shared_lock = shared_lock
self.verbose = verbose
self._stopping = threading.Event()
self.listen_thread = None
# connection does not connect to the server until either the
# connection.connect() method is called explicitly or until
# kombu calls it implicitly as needed.
self.connection = Connection(hostname=hostname,
userid=userid,
password=password,
virtual_host=virtual_host,
port=DEFAULT_SSL_PORT,
ssl=True)
build_exchange = Exchange(name=build_exchange_name, type='topic')
self.queues = [Queue(name='queue/%s/build' % userid,
exchange=build_exchange,
routing_key='build.#.finished',
durable=durable_queues,
auto_delete=not durable_queues)]
if treeherder_url:
jobaction_exchange = Exchange(name=jobaction_exchange_name, type='topic')
self.queues.append(Queue(name='queue/%s/jobactions' % userid,
exchange=jobaction_exchange,
routing_key='#',
durable=durable_queues,
auto_delete=not durable_queues))
def start(self):
"""Runs the `listen` method on a new thread."""
if self.listen_thread and self.listen_thread.is_alive():
logger.warning('AutophonePulseMonitor.start: listen thread already started')
return
logger.debug('AutophonePulseMonitor.start: listen thread starting')
self.listen_thread = threading.Thread(target=self.listen,
name='PulseMonitorThread')
self.listen_thread.daemon = True
self.listen_thread.start()
def stop(self):
"""Stops the pulse monitor listen thread."""
logger.debug('AutophonePulseMonitor stopping')
self._stopping.set()
self.listen_thread.join()
logger.debug('AutophonePulseMonitor stopped')
def is_alive(self):
return self.listen_thread.is_alive()
def listen(self):
logger.debug('AutophonePulseMonitor: start shared_lock.acquire')
self.shared_lock.acquire()
try:
consumer = self.connection.Consumer(self.queues,
callbacks=[self.handle_message],
accept=['json'],
auto_declare=False)
for queue in self.queues:
queue(self.connection).queue_declare(passive=False)
queue(self.connection).queue_bind()
with consumer:
while not self._stopping.is_set():
try:
logger.debug('AutophonePulseMonitor shared_lock.release')
self.shared_lock.release()
self.connection.drain_events(timeout=self.timeout)
except socket.timeout:
pass
except KeyboardInterrupt:
raise
finally:
logger.debug('AutophonePulseMonitor shared_lock.acquire')
self.shared_lock.acquire()
logger.debug('AutophonePulseMonitor.listen: stopping')
except:
logger.exception('AutophonePulseMonitor Exception')
finally:
logger.debug('AutophonePulseMonitor exit shared_lock.release')
self.shared_lock.release()
self.connection.release()
def handle_message(self, data, message):
if self._stopping.is_set():
return
message.ack()
if '_meta' in data and 'payload' in data:
self.handle_build(data, message)
if (self.treeherder_url and 'action' in data and
'project' in data and 'job_id' in data):
self.handle_jobaction(data, message)
def handle_build(self, data, message):
if self.verbose:
logger.debug(
'handle_build:\n'
'\tdata : %s\n'
'\tmessage: %s' % (
json.dumps(data, sort_keys=True, indent=4),
json.dumps(message.__dict__, sort_keys=True, indent=4)))
try:
build = data['payload']['build']
except (KeyError, TypeError), e:
logger.debug('AutophonePulseMonitor.handle_build_event: %s pulse build data' % e)
return
fields = (
'appName', # Fennec
'branch',
'buildid',
'comments',
'packageUrl',
'platform',
'robocopApkUrl',
'symbolsUrl',
'testsUrl',
'who'
)
required_fields = (
'appName', # Fennec
'branch', # mozilla-central, ...
'comments',
'packageUrl',
'platform', # android...
)
build_data = {}
builder_name = build['builderName']
build_data['builder_name'] = builder_name
build_data['build_type'] = 'debug' if 'debug' in builder_name else 'opt'
for property in build['properties']:
property_name = property[0]
if property_name in fields and len(property) > 1 and property[1]:
build_data[property_name] = type(property[1])(property[1])
for required_field in required_fields:
if required_field not in build_data or not build_data[required_field]:
return
if build_data['appName'] != 'Fennec':
return
if not build_data['platform'].startswith('android'):
return
if build_data['branch'] not in self.trees:
return
if build_data['platform'] not in self.platforms:
return
if build_data['build_type'] not in self.buildtypes:
return
if build_data['branch'] == 'try' and 'autophone' not in build_data['comments']:
return
self.build_callback(build_data)