def get_sentence_distances(self, channel: pika.channel.Channel,
method: pika.spec.Basic.Deliver,
properties: pika.spec.BasicProperties,
body: bytes):
response = {
"success": False,
"error": "",
"type": "SentenceDistances",
"distances":
[] # key-values: { "distance": np.float "metric": str "_id": str }
}
request = json.loads(body)
try:
if self.wait_matcher_is_ready(5):
response["distances"] = self.text_matcher.get_sorted_distances(
request["sentenceUUID"])
response["success"] = True
except Exception as err:
self.logger.error(err, exc_info=err)
response["error"] = str(err)
channel.basic_publish(exchange='' if properties.reply_to else "logs",
routing_key=properties.reply_to or "",
body=json.dumps(response),
properties=pika.BasicProperties(
correlation_id=properties.correlation_id,
content_type="application/json"))
channel.basic_ack(delivery_tag=method.delivery_tag)
def _log(self, channel: pika.channel.Channel, msg: Union[str, List[str]]):
log_data = {"Channel" : "Log",
"Node": self._task.node_id,
"Messages" : msg if isinstance(msg, list) else [msg]
}
log_body = json.dumps(log_data)
channel.basic_publish(exchange=self._pika.log_channel, routing_key='', body=log_body)
def _message_handler(self, channel: pika.channel.Channel,
method: pika.spec.Basic.Deliver,
properties: pika.spec.BasicProperties,
body: str) -> None:
message = json.loads(body)
print("Receiving: %s" % message)
self._work(message)
channel.basic_ack(method.delivery_tag)
def publish_message(queue_name: str, channel: pika.channel.Channel,
message: dict):
print('SENDING MESSAGE')
channel.basic_publish(exchange='',
properties=pika.BasicProperties(),
routing_key=queue_name,
body=json.dumps(message))
def add_on_channel_close_callback(self,
channel: pika.channel.Channel) -> None:
"""This method tells pika to call the on_channel_closed method if
RabbitMQ unexpectedly closes the channel.
"""
LOGGER.info("Adding channel close callback")
channel.add_on_close_callback(self.on_channel_closed)
def publish(self, connection: pika.BlockingConnection,
channel: pika.channel.Channel) -> None:
"""
Publishes a message to the channel
"""
kwargs = self.publish_kwargs()
kwargs['properties'] = pika.BasicProperties(**self.properties())
channel.basic_publish(**kwargs)
connection.close()
def on_channel_open(self, channel: pika.channel.Channel) -> None:
for backend, config in self.exchanges.items():
ex_name = config.get('exchange')
channel.exchange_declare(
ex_name,
exchange_type='topic',
durable=True,
)
logging.info(f'rabbitmq exchange: {ex_name} declared')
return
def on_channel(channel: pika.channel.Channel):
"""On channel closed handler.
"""
channel.add_on_close_callback(self.on_channel_closed)
channel.basic_qos(prefetch_count=100)
self._channels[name] = channel
try:
self._channels_opening[name].set_result(channel)
except asyncio.InvalidStateError: # pragma: no cover
pass
future.set_result(channel)
def _on_request_handler(channel: pika.channel.Channel, method: Any,
props: pika.BasicProperties, body: bytes) -> None:
print(f'{datetime.now()} > [x] Received requests. '
f'Message body: {body.decode(encoding="UTF-8")}.')
res: float = Main._compute(body.decode(encoding='UTF-8'))
channel.basic_publish(exchange='',
routing_key=props.reply_to,
properties=pika.BasicProperties(
correlation_id=props.correlation_id),
body=str(res))
channel.basic_ack(delivery_tag=method.delivery_tag)
print(f'{datetime.now()} > [x] Response sent.')
def test_pikathread_broadcast_reconnection(connection_params,
test_channel: pika.channel.Channel):
thread = _PikaThread(connection_params)
try:
thread.start()
thread.wait_for_connection()
got_message = threading.Event()
def _got_message(*args):
got_message.set()
exchange = test_channel.temporary_exchange_declare(
exchange_type="fanout")
thread.subscribe_broadcast(exchange,
_got_message,
reconnectable=True,
subscription_id=1).result()
# Force reconnection - normally we want this to be transparent, but
# let's twiddle the internals so we can wait for reconnection as we
# don't want to pick up the message before it resets
print("Terminating connection")
thread._connected.clear()
thread._debug_close_connection()
thread.wait_for_connection()
print("Reconnected")
# Now, make sure we still get this message
test_channel.basic_publish(exchange, routing_key="", body="A Message")
got_message.wait(2)
assert got_message.is_set()
# Add a non-resubscribable connection
got_message_2 = threading.Event()
def _got_message_2(*args):
got_message_2.set()
thread.subscribe_broadcast(
exchange,
_got_message_2,
reconnectable=False,
subscription_id=2,
).result()
# Make sure that the thread ends instead of reconnect if we force a disconnection
thread._debug_close_connection()
thread.join()
finally:
thread.join(stop=True)
def declare_queues(channel: pika.channel.Channel, queues: typing.List):
"""Declare queues
:params: channel Channel of RabbitMQ connections
:params: queues List of queues need declaring
"""
logger.info('[BEGIN] declaring queues')
for q in queues:
channel.queue_declare(q['name'],
durable=q['durable'],
auto_delete=q['auto_delete'],
exclusive=q['exclusive'])
logger.info(f"- Queue {q['name']}: done")
logger.info('[END] declaring queues')
def declare_exchanges(channel: pika.channel.Channel, exchanges: typing.List):
"""Declare exchanges
:params: channel Channel of RabbitMQ connections
:params: exchanges List of exchanges need declaring
"""
logger.info('[BEGIN] declaring exchanges')
for e in exchanges:
channel.exchange_declare(
e['name'],
e_type=e['type'],
durable=e['durable'],
auto_delete=e['auto_delete'],
internal=e['internal'],
)
logger.info(f"- Exchange {e['name']}: done")
logger.info('[END] declaring exchanges')
def reply_sync(
self,
ch: pika.channel.Channel,
method: pika.spec.Basic.Deliver,
properties: pika.spec.BasicProperties,
body: str,
app=None,
):
"""
Will reply to a message, which was send by :meth:`publish_sync`
Example::
@queue(queue_name="rpc")
def concat_callback(ch, method, props, body):
result = body["a"] + body["b"]
body = {"result": result}
coney.reply_sync(ch, method, props, body)
This is a conveniences short hand method for::
@queue(queue_name="rpc")
def concat_callback(ch, method, props, body):
result = body["a"] + body["b"]
body = {"result": result}
self.publish(
body,
routing_key=properties.reply_to,
properties={"correlation_id": properties.correlation_id},
app=app,
)
:parameter ch:
:parameter method:
:parameter properties:
:parameter body: The message to send
"""
self.publish(
body,
routing_key=properties.reply_to,
properties={"correlation_id": properties.correlation_id},
app=app,
)
ch.basic_ack(delivery_tag=method.delivery_tag)
def bind(channel: pika.channel.Channel, bindings: typing.List):
"""Binding from source to destination
:params: channel Channel of RabbitMQ connections
:params: bindings List of bindings need to be bound
"""
logger.info('[BEGIN] binding')
for b in bindings:
if b['destination_type'] == 'queue':
channel.queue_bind(queue=b['destination'],
exchange=b['source'],
routing_key=b['routing_key'])
elif b['destination_type'] == 'exchange':
channel.exchange_bind(destination=b['destination'],
source=b['source'],
routing_key=b['routing_key'])
logger.info(
f"- Binding {b['destination_type']} \"{b['destination']}\" to \"{b['source']}\" by \"{b['routing_key']}\""
)
logger.info('[END] binding')
def setup_exchange(self, exchange_name: str, exchange_type: str,
channel: pika.channel.Channel) -> None:
"""Setup the exchange on RabbitMQ by invoking the Exchange.Declare RPC
command. When it is complete, the on_exchange_declareok method will
be invoked by pika.
:param str|unicode exchange_name: The name of the exchange to declare
"""
LOGGER.info("Declaring exchange: %s", exchange_name)
# Note: using functools.partial is not required, it is demonstrating
# how arbitrary data can be passed to the callback when it is called
cb = functools.partial(self.on_exchange_declareok,
exchange_name=exchange_name)
channel.exchange_declare(
exchange=exchange_name,
exchange_type=exchange_type,
callback=cb,
durable=True,
)
def ack_policy(
self,
ch: pika.channel.Channel,
deliver: pika.spec.Basic.Deliver,
conversation_id: str,
correlation_id: str,
) -> Generator[None, None, None]:
try:
yield None
except Exception:
error = traceback.format_exc()
self.log(
ALERT,
f"Unhandled error while processing message {deliver.routing_key}",
error,
conversation_id=conversation_id,
)
# retry once
if not deliver.redelivered:
ch.basic_nack(delivery_tag=deliver.delivery_tag, requeue=True)
else:
# dead letter
self.log(
EMERGENCY,
f"Giving up on {deliver.routing_key}",
error,
conversation_id=conversation_id,
)
ch.basic_nack(delivery_tag=deliver.delivery_tag, requeue=False)
else:
ch.basic_ack(delivery_tag=deliver.delivery_tag)
def on_message(self, _unused_channel: pika.channel.Channel, basic_deliver: Basic.Deliver,
properties: BasicProperties, body: bytes):
"""Invoked by pika when a message is delivered from RabbitMQ. The
channel is passed for your convenience. The basic_deliver object that
is passed in carries the exchange, routing key, delivery tag and
a redelivered flag for the message. The properties passed in is an
instance of BasicProperties with the message properties and the body
is the message that was sent.
:param pika.channel.Channel _unused_channel: The channel object
:param pika.Spec.Basic.Deliver basic_deliver: basic_deliver method
:param pika.Spec.BasicProperties properties: properties
:param bytes body: The message body
"""
self._logger.info('Received message # %s from %s: %s',
basic_deliver.delivery_tag, properties.app_id, body)
try:
data = json.loads(body)
except Exception as exc:
self._logger.exception(f"exception while decode data {exc}")
_unused_channel.basic_reject(basic_deliver.delivery_tag, requeue=False)
else:
self._callback(_unused_channel, basic_deliver, properties, data)
self.acknowledge_message(basic_deliver.delivery_tag)
def new_text(self, channel: pika.channel.Channel,
method: pika.spec.Basic.Deliver,
properties: pika.spec.BasicProperties, body: bytes):
request = json.loads(body)
response = {
"success": False,
"error": "",
"textUUID": "",
"type": "processText"
}
if self.wait_matcher_is_ready(5):
try:
response["textUUID"] = self.text_matcher.process_new_text(
request)
response["success"] = True
except Exception as err:
self.logger.error(err, exc_info=err)
response["error"] = str(err)
else:
response[
"error"] = f"TextMatcher is not IDLE: {self.text_matcher.state}"
channel.basic_publish(exchange='',
routing_key=properties.reply_to,
body=json.dumps(response),
properties=pika.BasicProperties(
correlation_id=properties.correlation_id,
content_type="application/json"))
channel.basic_ack(delivery_tag=method.delivery_tag)
channel.basic_publish(
exchange="logs",
routing_key="",
body=json.dumps({
"level":
"info",
"type":
"success",
"message":
f"Successfully processed text (ID: {response['textUUID']})."
}),
properties=pika.BasicProperties(content_type="application/json"))
def declare_amqp_pipeline(conf: config,
channel: pika.channel.Channel,
durable: bool = True) -> None:
"""Declare AMQP Pipeline.
This function declares all necessary exchanges and queues based on conf.project_name aka. does the plumbing
"""
# agents to collector
channel.exchange_declare(exchange=conf.name_exchange_agents,
exchange_type='fanout')
queue_agents = channel.queue_declare(queue=conf.name_queue_agents,
durable=durable)
channel.queue_bind(exchange=conf.name_exchange_agents,
queue=queue_agents.method.queue)
# collector to analysers
channel.exchange_declare(exchange=conf.name_exchange_analyser,
exchange_type='fanout')
queue_analyser_addr = channel.queue_declare(
queue=conf.name_queue_analyser_addr, durable=durable)
queue_analyser_entropy = channel.queue_declare(
queue=conf.name_queue_analyser_entropy, durable=durable)
queue_analyser_lof = channel.queue_declare(
queue=conf.name_queue_analyser_lof, durable=durable)
queue_analyser_svm = channel.queue_declare(
queue=conf.name_queue_analyser_svm, durable=durable)
channel.queue_bind(exchange=conf.name_exchange_analyser,
queue=queue_analyser_addr.method.queue)
channel.queue_bind(exchange=conf.name_exchange_analyser,
queue=queue_analyser_entropy.method.queue)
channel.queue_bind(exchange=conf.name_exchange_analyser,
queue=queue_analyser_lof.method.queue)
channel.queue_bind(exchange=conf.name_exchange_analyser,
queue=queue_analyser_svm.method.queue)
# only 1 packet to process at a time
channel.basic_qos(prefetch_count=1)
def on_message(
self,
ch: pika.channel.Channel,
basic_deliver: pika.spec.Basic.Deliver,
properties: pika.spec.BasicProperties,
body: bytes,
) -> None:
"""Invoked by pika when a message is delivered from RabbitMQ. The
channel is passed for your convenience. The basic_deliver object that
is passed in carries the exchange, routing key, delivery tag and
a redelivered flag for the message. The properties passed in is an
instance of BasicProperties with the message properties and the body
is the message that was sent.
:param pika.channel.Channel ch: The channel object
:param pika.Spec.Basic.Deliver: basic_deliver method
:param pika.Spec.BasicProperties: properties
:param bytes body: The message body
"""
headers: HEADER = properties.headers
decoder = cbor if properties.content_type == "application/cbor" else json
routing_key: str = basic_deliver.routing_key
exchange: str = basic_deliver.exchange
if headers is None or "conversation_id" not in headers:
self.log(EMERGENCY, f"Missing headers on {routing_key}")
# unrecoverable error, send to dead letter
ch.basic_nack(delivery_tag=basic_deliver.delivery_tag,
requeue=False)
return
conversation_id = headers.pop("conversation_id")
try:
payload: JSON_MODEL = decoder.loads(body) if body else None
except ValueError:
self.log(
EMERGENCY,
f"Unable to decode payload for {routing_key}; dead lettering.",
conversation_id=conversation_id,
)
# Unrecoverable, put to dead letter
ch.basic_nack(delivery_tag=basic_deliver.delivery_tag,
requeue=False)
return
LOGGER.info("Received message from %s: %s", exchange, routing_key)
# meta
headers["timestamp"] = properties.timestamp
headers["expiration"] = properties.expiration
headers["user_id"] = properties.user_id
headers["app_id"] = properties.app_id
if exchange == self.CMD_EXCHANGE:
correlation_id = properties.correlation_id
reply_to = properties.reply_to
status = headers.pop("status", "") if headers else ""
if not (reply_to or status):
self.log(
EMERGENCY,
"invalid enveloppe for command/result: {}; dead lettering."
.format(headers),
conversation_id=conversation_id,
)
# Unrecoverable. Put to dead letter
ch.basic_nack(delivery_tag=basic_deliver.delivery_tag,
requeue=False)
return
if reply_to:
with self.ack_policy(ch, basic_deliver, conversation_id,
correlation_id):
self.handle_command(
routing_key,
payload,
conversation_id,
reply_to,
correlation_id,
meta=headers,
)
else:
with self.ack_policy(ch, basic_deliver, conversation_id,
correlation_id):
self.handle_result(
routing_key,
payload,
conversation_id,
status,
correlation_id,
meta=headers,
)
else:
with self.ack_policy(ch, basic_deliver, conversation_id, ""):
self.handle_event(routing_key,
payload,
conversation_id,
meta=headers)