mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 09:36:30 +00:00
Alex Bezdieniezhnykh
116 lines
4.3 KiB
Cython
116 lines
4.3 KiB
Cython
import asyncio
|
|
import json
|
|
import socket
|
|
import struct
|
|
import threading
|
|
|
|
import msgpack
|
|
from msgpack import packb
|
|
from rstream import Producer, Consumer, AMQPMessage, ConsumerOffsetSpecification, OffsetType, MessageContext
|
|
|
|
cimport constants
|
|
from api_client cimport ApiClient
|
|
from processor_command cimport FileCommand, ProcessorType
|
|
from annotation cimport Annotation
|
|
|
|
cdef class QueueConfig:
|
|
cdef str host,
|
|
cdef int port
|
|
cdef str producer_user, producer_pw, consumer_user, consumer_pw
|
|
|
|
@staticmethod
|
|
cdef QueueConfig from_json(str json_string):
|
|
cdef dict config_dict = json.loads(<str>json_string)["QueueConfig"]
|
|
cdef QueueConfig config = QueueConfig()
|
|
|
|
config.host = config_dict["Host"]
|
|
config.port = config_dict["Port"]
|
|
config.producer_user = config_dict["ProducerUsername"]
|
|
config.producer_pw = config_dict["ProducerPassword"]
|
|
config.consumer_user = config_dict["ConsumerUsername"]
|
|
config.consumer_pw = config_dict["ConsumerPassword"]
|
|
return config
|
|
|
|
cdef class SocketHandler:
|
|
"""Handles socket communication with size-prefixed messages."""
|
|
|
|
def __init__(self, object on_message):
|
|
self.on_message = on_message
|
|
self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
self._connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
self._socket.bind((constants.SOCKET_HOST, constants.SOCKET_PORT))
|
|
self._socket.listen(1)
|
|
|
|
cdef start(self):
|
|
threading.Thread(target=self.start_inner, daemon=True).start()
|
|
|
|
cdef start_inner(self):
|
|
while True:
|
|
self._connection, client_address = self._socket.accept()
|
|
size_data = self._connection.recv(4)
|
|
if not size_data:
|
|
raise ConnectionError("Connection closed while reading size prefix.")
|
|
data_size = struct.unpack('>I', size_data)[0]
|
|
# Read the full message
|
|
data = b""
|
|
while len(data) < data_size:
|
|
packet = self._socket.recv(data_size - len(data))
|
|
if not packet:
|
|
raise ConnectionError("Connection closed while reading data.")
|
|
data += packet
|
|
cmd = FileCommand.from_msgpack(data, ProcessorType.SOCKET)
|
|
self.on_message(cmd)
|
|
|
|
cdef send(self, list[Annotation] message):
|
|
data = msgpack.packb(message)
|
|
size_prefix = len(data).to_bytes(4, 'big')
|
|
self._connection.sendall(size_prefix + data)
|
|
|
|
cdef close(self):
|
|
if self._socket:
|
|
self._socket.close()
|
|
self._socket = None
|
|
|
|
cdef class RabbitHandler:
|
|
def __init__(self, ApiClient api_client, object on_message):
|
|
self.on_message = on_message
|
|
cdef str config_str = api_client.load_queue_config()
|
|
queue_config = QueueConfig.from_json(config_str)
|
|
self.annotation_producer = Producer(
|
|
host=<str>queue_config.host,
|
|
port=queue_config.port,
|
|
username=<str>queue_config.producer_user,
|
|
password=<str>queue_config.producer_pw
|
|
)
|
|
self.command_consumer = Consumer(
|
|
host=<str>queue_config.host,
|
|
port=queue_config.port,
|
|
username=<str>queue_config.consumer_user,
|
|
password=<str>queue_config.consumer_pw
|
|
)
|
|
|
|
cdef start(self):
|
|
threading.Thread(target=self._run_async, daemon=True).start()
|
|
|
|
def _run_async(self):
|
|
asyncio.run(self.start_inner())
|
|
|
|
async def start_inner(self):
|
|
await self.command_consumer.start()
|
|
await self.command_consumer.subscribe(stream=<str>constants.COMMANDS_QUEUE, callback=self.on_message_inner,
|
|
offset_specification=ConsumerOffsetSpecification(OffsetType.FIRST, None)) # put real offset
|
|
|
|
def on_message_inner(self, message: AMQPMessage, message_context: MessageContext):
|
|
cdef bytes body = message.body
|
|
cmd = FileCommand.from_msgpack(body, ProcessorType.RABBIT)
|
|
self.on_message(cmd)
|
|
|
|
cdef send(self, object message):
|
|
packed_message = AMQPMessage(body=packb(message))
|
|
self.annotation_producer.send(<str>constants.ANNOTATIONS_QUEUE, packed_message)
|
|
|
|
cdef close(self):
|
|
if self.annotation_producer:
|
|
self.annotation_producer.close()
|
|
if self.command_consumer:
|
|
self.command_consumer.close() |