Add support for file streaming upload, update to async pika lib #9
@@ -1,4 +1,5 @@
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
from io import BytesIO
|
||||
from datetime import datetime, timezone
|
||||
@@ -8,6 +9,7 @@ from amqp.model.model import AMQMessage
|
||||
from amqp.model.snowflake_id import SnowflakeId
|
||||
from amqp.router.serializer import map_with_list_as_string, map_as_string, parse_map_list_string, parse_map_string
|
||||
from amqp.router.utils import sanitize_as_rabbitmq_name
|
||||
from amqp.adapter.file_downloader import download_buffer
|
||||
|
||||
|
||||
class AMQMessageFactory:
|
||||
@@ -206,3 +208,20 @@ class AMQMessageFactory:
|
||||
trace_info,
|
||||
body_b64,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def from_stream(stream: bytes, rabbitmq_url: str, loop) -> AMQMessage | None:
|
||||
"""
|
||||
Builds the AMQMessage from its serialized (byte array) form.
|
||||
:param stream: serialized message
|
||||
:param rabbitmq_url: RabbitMQ connection URL
|
||||
:return: AMQMessage instance
|
||||
"""
|
||||
# async implementation
|
||||
message_data = json.loads(stream.decode())
|
||||
req_info = message_data.get('req-info', '')
|
||||
if req_info:
|
||||
(body, eof) = await download_buffer(loop=loop, rabbitmq_url=rabbitmq_url, queue_name=req_info)
|
||||
if eof == 1:
|
||||
return AMQMessageFactory.from_bytes(body)
|
||||
return None
|
||||
|
||||
@@ -1,38 +1,27 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import python_multipart
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import threading
|
||||
import traceback
|
||||
|
||||
from aiohttp import ClientSession, ClientResponse
|
||||
from fastapi import HTTPException, UploadFile
|
||||
from python_multipart.multipart import Field, File
|
||||
from fastapi import HTTPException
|
||||
from typing import Dict, List, Tuple, Callable, Coroutine, AnyStr, Any
|
||||
from urllib.parse import parse_qs
|
||||
from xml.etree import ElementTree
|
||||
|
||||
from amqp.adapter.amq_response_factory import AMQResponseFactory
|
||||
from amqp.service.multipart_parser import CleverMultiPartParser
|
||||
from amqp.config.amq_configuration import AMQConfiguration
|
||||
from amqp.model.model import AMQMessage, AMQResponse
|
||||
|
||||
|
||||
class CleverMultiPartParser:
|
||||
def __init__(self, message: AMQMessage):
|
||||
self.form_data = {}
|
||||
# Convert each list of strings to a single string joined by newline, then to bytes
|
||||
headers = {key: "\n".join(value).encode("utf-8") for key, value in message.headers.items()}
|
||||
python_multipart.parse_form(headers, io.BytesIO(message.body()), self.on_field, self.on_file)
|
||||
|
||||
def on_field(self, field: Field):
|
||||
print(field)
|
||||
self.form_data[field.field_name.decode('utf-8')] = field.value
|
||||
print(self.form_data)
|
||||
|
||||
def on_file(self, file: File):
|
||||
print(file)
|
||||
self.form_data[file.field_name.decode('utf-8')] = \
|
||||
UploadFile(file.file_object, size=file.size, filename=file.file_name)
|
||||
print(self.form_data)
|
||||
# Track the number of messages currently being processed
|
||||
current_parallel_executions = 0
|
||||
lock = threading.Lock()
|
||||
parallel_workers = os.getenv("PARALLEL_WORKERS", default=10)
|
||||
|
||||
|
||||
def parse_request_body(message: AMQMessage):
|
||||
@@ -45,15 +34,17 @@ def parse_request_body(message: AMQMessage):
|
||||
Returns:
|
||||
dict: Parsed data as a dictionary.
|
||||
"""
|
||||
content_type = message.headers.get('Content-Type', '')
|
||||
if not content_type:
|
||||
_content_type = message.headers.get('Content-Type', '')
|
||||
if not _content_type:
|
||||
raise ValueError("Content-Type header is required")
|
||||
|
||||
content_type = _content_type[0]
|
||||
# Parse JSON
|
||||
if content_type == "application/json":
|
||||
try:
|
||||
return json.loads(message.body())
|
||||
except json.JSONDecodeError as e:
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
|
||||
raise ValueError(f"Invalid JSON: {e}")
|
||||
|
||||
# Parse URL-encoded form data
|
||||
@@ -61,6 +52,8 @@ def parse_request_body(message: AMQMessage):
|
||||
try:
|
||||
return {k: v[0] if len(v) == 1 else v for k, v in parse_qs(message.body()).items()}
|
||||
except Exception as e:
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
|
||||
raise ValueError(f"Invalid URL-encoded form data: {e}")
|
||||
|
||||
# Parse multipart/form-data
|
||||
@@ -69,6 +62,8 @@ def parse_request_body(message: AMQMessage):
|
||||
parser: CleverMultiPartParser = CleverMultiPartParser(message)
|
||||
return parser.form_data
|
||||
except Exception as e:
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
|
||||
raise ValueError(f"Invalid multipart/form-data: {e}")
|
||||
|
||||
# Parse plain text
|
||||
@@ -113,9 +108,13 @@ async def call_cleverthis_post_put_api(
|
||||
result.encode('utf-8'), message.trace_info
|
||||
)
|
||||
except HTTPException as http_error:
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {http_error}")
|
||||
return AMQResponseFactory.of(message.id, http_error.status_code, "text/plain",
|
||||
str(http_error.detail).encode('utf-8'), message.trace_info)
|
||||
except Exception as error:
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {error}")
|
||||
return AMQResponseFactory.of(message.id, 500, "text/plain", str(error).encode('utf-8'), message.trace_info)
|
||||
|
||||
|
||||
@@ -144,9 +143,13 @@ async def call_cleverthis_get_api(
|
||||
_msg = f"Unexpected URL: {message.path}, expected pattern: {pattern}"
|
||||
return AMQResponseFactory.of(message.id, 404, "text/plain", _msg.encode('utf-8'), message.trace_info)
|
||||
except HTTPException as http_error:
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {http_error}")
|
||||
return AMQResponseFactory.of(message.id, http_error.status_code, "text/plain",
|
||||
str(http_error.detail).encode('utf-8'), message.trace_info)
|
||||
except Exception as error:
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {error}")
|
||||
return AMQResponseFactory.of(message.id, 500, "text/plain", str(error).encode('utf-8'), message.trace_info)
|
||||
|
||||
|
||||
@@ -170,6 +173,14 @@ class CleverThisServiceAdapter:
|
||||
"""
|
||||
return None
|
||||
|
||||
def sync_on_message(self, message: AMQMessage) -> AMQResponse:
|
||||
"""
|
||||
Synchronous version of on_message method. This is a wrapper around the async on_message method.
|
||||
:param message: AMQMessage
|
||||
:return: AMQResponse
|
||||
"""
|
||||
return asyncio.run(self.on_message(message))
|
||||
|
||||
async def on_message(self, message) -> AMQResponse:
|
||||
"""
|
||||
Called when a message is received from the AMQP.
|
||||
@@ -177,30 +188,46 @@ class CleverThisServiceAdapter:
|
||||
:return: AMQP response class with operation status code and optional error details.
|
||||
"""
|
||||
DEBUG_NO_AUTH = 1
|
||||
global current_parallel_executions
|
||||
|
||||
# Increment the counter for parallel executions
|
||||
with lock:
|
||||
current_parallel_executions += 1
|
||||
logging.info(f"PARALLEL: Current Capacity [{current_parallel_executions}]")
|
||||
if current_parallel_executions > parallel_workers - 2:
|
||||
logging.warning("Warning: Capacity close to depleted!")
|
||||
|
||||
_auth: str = message.headers.get('Authorization', '')
|
||||
_auth_user: Any | None = None
|
||||
if _auth == '' and DEBUG_NO_AUTH == 1:
|
||||
_auth = f"Bearer DEBUG_NO_AUTH"
|
||||
|
||||
if 'Bearer' in _auth:
|
||||
try:
|
||||
token = _auth.split(' ')[1]
|
||||
_auth_user = await self.get_current_user(token)
|
||||
except Exception as error:
|
||||
logging.error(f"Error getting user: {error}")
|
||||
amq_response: AMQResponse | None = None
|
||||
if _auth_user or DEBUG_NO_AUTH == 1:
|
||||
method = message.method.lower()
|
||||
if method == 'get':
|
||||
return await self.get(message, _auth_user)
|
||||
amq_response = await self.get(message, _auth_user)
|
||||
elif method == 'put':
|
||||
return await self.put(message, _auth_user)
|
||||
amq_response = await self.put(message, _auth_user)
|
||||
elif method == 'post':
|
||||
return await self.post(message, _auth_user)
|
||||
amq_response = await self.post(message, _auth_user)
|
||||
elif method == 'head':
|
||||
return await self.head(message, _auth_user)
|
||||
amq_response = await self.head(message, _auth_user)
|
||||
elif method == 'delete':
|
||||
return await self.delete(message, _auth_user)
|
||||
amq_response = await self.delete(message, _auth_user)
|
||||
else:
|
||||
raise ValueError(f"Unexpected HTTP method: {message.method}")
|
||||
else:
|
||||
return AMQResponseFactory.of(message.id, 401, "text/plain", b"Unauthorized", message.trace_info)
|
||||
amq_response = AMQResponseFactory.of(message.id, 401, "text/plain", b"Unauthorized", message.trace_info)
|
||||
with lock:
|
||||
current_parallel_executions -= 1
|
||||
return amq_response
|
||||
|
||||
# ==================================================================================================
|
||||
# ===================== C l e v e r T h i s A P I m e t h o d s ============================
|
||||
@@ -236,6 +263,7 @@ class CleverThisServiceAdapter:
|
||||
:return: AMQResponse to be sent back to the API Gateway via AMQP
|
||||
"""
|
||||
return await self.handle_possible_form(
|
||||
message,
|
||||
self.session.put(
|
||||
f"http://{self.service_host}:{self.service_port}{message.path}",
|
||||
data=message.body,
|
||||
@@ -253,6 +281,7 @@ class CleverThisServiceAdapter:
|
||||
:return: AMQResponse to be sent back to the API Gateway via AMQP
|
||||
"""
|
||||
return await self.handle_possible_form(
|
||||
message,
|
||||
self.session.post(
|
||||
f"http://{self.service_host}:{self.service_port}{message.path}",
|
||||
data=message.body,
|
||||
@@ -261,7 +290,8 @@ class CleverThisServiceAdapter:
|
||||
auth_user
|
||||
)
|
||||
|
||||
async def handle_possible_form(self, request_coroutine, auth_user: Any) -> AMQResponse:
|
||||
async def handle_possible_form(self, message: AMQMessage, request_coroutine, auth_user: Any) -> AMQResponse:
|
||||
args = parse_request_body(message)
|
||||
return await request_coroutine
|
||||
|
||||
async def head(self, message: AMQMessage, auth_user: Any) -> AMQResponse:
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from aio_pika import connect_robust
|
||||
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustChannel, AbstractRobustQueue, AbstractQueue
|
||||
|
||||
# Global dictionary to track completed downloads
|
||||
download_locations = defaultdict(str)
|
||||
IN_PROGRESS = 0
|
||||
END_OF_FILE = 1
|
||||
CANCELLED = 2
|
||||
|
||||
|
||||
def ensure_directory_exists(filepath):
|
||||
"""Ensures that the directory for the given filepath exists."""
|
||||
dir_name = os.path.dirname(filepath)
|
||||
if dir_name and not os.path.exists(dir_name):
|
||||
os.makedirs(dir_name)
|
||||
|
||||
|
||||
async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir) -> (int, int):
|
||||
"""
|
||||
Asynchronously downloads a file from RabbitMQ based on the provided file information.
|
||||
takes the file_info and message_data and output_dir, creates queue consumer and waits for the file to be downloaded
|
||||
Args:
|
||||
loop: The asyncio event loop.
|
||||
rabbitmq_url: The RabbitMQ connection URL.
|
||||
file_info (dict): A dictionary containing file metadata
|
||||
message_data (dict): The entire message
|
||||
output_dir (str): The directory where the file should be saved.
|
||||
"""
|
||||
global download_locations
|
||||
filename = file_info['filename']
|
||||
size = file_info['size']
|
||||
stream_name = file_info['streamName']
|
||||
queue_name = file_info['queue_name']
|
||||
filepath = os.path.join(output_dir, filename)
|
||||
ensure_directory_exists(filepath)
|
||||
download_locations[file_info['key']] = filepath
|
||||
eof = IN_PROGRESS
|
||||
fsize = 0
|
||||
logging.info(f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}")
|
||||
if queue_name:
|
||||
connection = await connect_robust(rabbitmq_url, loop=loop)
|
||||
try:
|
||||
channel: AbstractRobustChannel = await connection.channel()
|
||||
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
|
||||
logging.info(f"Awaiting all file chunks being downloaded for {filename}")
|
||||
async with queue.iterator() as queue_iter:
|
||||
async for message in queue_iter:
|
||||
async with message.process():
|
||||
with open(filepath, 'ab') as f:
|
||||
f.write(message.body)
|
||||
fsize = fsize + message.body_size
|
||||
eof = message.headers.get('eof', END_OF_FILE)
|
||||
if eof != IN_PROGRESS:
|
||||
break
|
||||
except Exception as e:
|
||||
logging.error(f"Error downloading file: {e}")
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
|
||||
await connection.close()
|
||||
return (fsize, eof)
|
||||
|
||||
|
||||
async def download_buffer(loop, rabbitmq_url, queue_name, ttl = 604800000) -> (bytearray, int):
|
||||
"""
|
||||
Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
|
||||
It assumes entire array is in the first (and only) chunk
|
||||
Args:
|
||||
loop: The asyncio event loop.
|
||||
rabbitmq_url: The RabbitMQ connection URL.
|
||||
queue_name (str): The directory where the file should be saved.
|
||||
"""
|
||||
res = bytearray()
|
||||
eof = IN_PROGRESS
|
||||
logging.info(f"Downloading buffer from stream: {queue_name}")
|
||||
if queue_name:
|
||||
connection = await connect_robust(rabbitmq_url, loop=loop)
|
||||
try:
|
||||
channel: AbstractRobustChannel = await connection.channel()
|
||||
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
|
||||
logging.info(f"Awaiting all buffer chunks being downloaded for {queue_name}")
|
||||
async with queue.iterator() as queue_iter:
|
||||
async for message in queue_iter:
|
||||
async with message.process():
|
||||
res.extend(message.body)
|
||||
eof = message.headers.get('eof', END_OF_FILE)
|
||||
if eof != IN_PROGRESS:
|
||||
break
|
||||
logging.info(f"Finished downloading buffer {queue_name}")
|
||||
except Exception as e:
|
||||
logging.error(f"Error downloading buffer: {e}")
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
|
||||
await connection.close()
|
||||
return (res, eof)
|
||||
|
||||
|
||||
async def on_message(
|
||||
message: AbstractIncomingMessage,
|
||||
json_data: bytes,
|
||||
rabbitmq_url: str,
|
||||
loop,
|
||||
output_dir="downloaded_files"
|
||||
) -> defaultdict:
|
||||
"""
|
||||
Asynchronous callback function for handling messages from the 'amq-cleverswarm' queue.
|
||||
|
||||
Args:
|
||||
message: The incoming RabbitMQ message.
|
||||
rabbitmq_url: The RabbitMQ connection URL.
|
||||
"""
|
||||
global download_locations
|
||||
try:
|
||||
message_data = json.loads(message.body.decode())
|
||||
amq_message_data = json.loads(json_data.decode())
|
||||
|
||||
files_to_download = []
|
||||
for key, file_type in amq_message_data['files'].items():
|
||||
if len(file_type) > 0:
|
||||
files_info = file_type[0]
|
||||
stream_name = files_info['streamName']
|
||||
files_info['queue_name'] = message_data.get(stream_name, "") if stream_name else ""
|
||||
files_info['key'] = key
|
||||
files_to_download.append(files_info)
|
||||
|
||||
if not files_to_download:
|
||||
logging.info("No files to download in this message.")
|
||||
await message.ack()
|
||||
return False
|
||||
|
||||
# Create a task for each file download
|
||||
tasks = []
|
||||
for file_info in files_to_download:
|
||||
logging.info(f"about ti create task for {file_info}")
|
||||
task = asyncio.create_task(download_file(loop, rabbitmq_url, file_info, message_data, output_dir))
|
||||
tasks.append(task)
|
||||
logging.info("Waiting for all files to download...gather(*tasks)")
|
||||
await asyncio.gather(*tasks) # Wait for all downloads to complete
|
||||
|
||||
# After all files are downloaded, acknowledge the original message
|
||||
logging.info(f"All files downloaded. Acknowledging the main message d-tag: {message.delivery_tag}")
|
||||
await message.ack()
|
||||
logging.info(f"Message d-tag: {message.delivery_tag} acknowledged.")
|
||||
except Exception as e:
|
||||
logging.error(f"[E] [{message.delivery_tag}] Error processing message: {e}")
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f" Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
|
||||
await message.ack() # ignore the message
|
||||
|
||||
return download_locations
|
||||
@@ -0,0 +1,140 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
|
||||
from aio_pika.abc import HeadersType, AbstractIncomingMessage, AbstractQueue, AbstractRobustChannel
|
||||
from aiormq import Channel
|
||||
|
||||
|
||||
class MultipartDownload:
|
||||
|
||||
def __init__(self, headers: HeadersType, channel: AbstractRobustChannel):
|
||||
self.headers = headers
|
||||
self.channel = channel
|
||||
self.addressing: int = headers.get('clevermicro.addressing', 0)
|
||||
self.files_to_download = []
|
||||
self.threads = []
|
||||
# Global dictionary to track completed downloads
|
||||
self.download_status = defaultdict(lambda: defaultdict(bool))
|
||||
self.future = None
|
||||
|
||||
def ensure_directory_exists(self, filepath):
|
||||
"""Ensures that the directory for the given filepath exists."""
|
||||
dir_name = os.path.dirname(filepath)
|
||||
if dir_name and not os.path.exists(dir_name):
|
||||
os.makedirs(dir_name)
|
||||
|
||||
async def start_queue2(self, channel: Channel, queue_name: str, callback):
|
||||
dok = await channel.queue_declare(queue=queue_name, durable=False, auto_delete=True)
|
||||
await channel.basic_consume(queue=queue_name, consumer_callback=callback, no_ack=True, exclusive=True)
|
||||
|
||||
async def start_queue(self, queue_name: str, callback):
|
||||
"""
|
||||
Starts consuming from a queue.
|
||||
|
||||
Args:
|
||||
:param queue_name: The name of the queue to consume from.
|
||||
:param callback: The callback function to call when a message is received.
|
||||
"""
|
||||
queue: AbstractQueue = await self.channel.declare_queue(name=queue_name, durable=False, auto_delete=True)
|
||||
await queue.consume(callback=callback, no_ack=True, exclusive=True)
|
||||
|
||||
def download_file(self, channel: Channel, file_info, message_data, output_dir):
|
||||
"""
|
||||
Downloads a file from RabbitMQ based on the provided file information.
|
||||
|
||||
Args:
|
||||
channel: The RabbitMQ channel.
|
||||
file_info (dict): A dictionary containing file metadata
|
||||
message_data (dict): the entire message
|
||||
output_dir (str): The directory where the file should be saved.
|
||||
"""
|
||||
filename = file_info['filename']
|
||||
size = file_info['size']
|
||||
stream_name = file_info['streamName']
|
||||
filepath = os.path.join(output_dir, filename)
|
||||
self.ensure_directory_exists(filepath)
|
||||
|
||||
print(f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name} to {filepath}")
|
||||
|
||||
bytes_received = 0
|
||||
with open(filepath, 'wb') as f:
|
||||
async def callback(in_message: AbstractIncomingMessage):
|
||||
nonlocal bytes_received
|
||||
f.write(in_message.body)
|
||||
bytes_received += len(in_message.body)
|
||||
# print(f"Received {len(body)} bytes for {filename}. Total: {bytes_received}/{size}") # uncomment for verbose
|
||||
|
||||
if bytes_received >= size:
|
||||
print(f"Finished downloading {filename}")
|
||||
f.close()
|
||||
# Update global download status
|
||||
message_req_info = json.loads(message_data.get('req-info', '{}'))
|
||||
if message_req_info:
|
||||
self.download_status[message_req_info['from']][filename] = True
|
||||
await in_message.channel.close() # stop consuming after the file is downloaded
|
||||
|
||||
self.start_queue2(channel, stream_name, callback)
|
||||
|
||||
def on_message(self, in_message: AbstractIncomingMessage, body: bytes) -> asyncio.Future:
|
||||
"""
|
||||
Callback function for handling messages from the 'amq-cleverswarm' queue.
|
||||
|
||||
Args:
|
||||
ch: The RabbitMQ channel.
|
||||
method: The delivery method.
|
||||
properties: The message properties.
|
||||
body: The message body (a JSON string).
|
||||
:param body:
|
||||
"""
|
||||
message_data = json.loads(body.decode())
|
||||
req_info = json.loads(message_data.get('req-info', '{}')) # Handle missing 'req-info'
|
||||
self.future = asyncio.Future()
|
||||
|
||||
if not req_info:
|
||||
print("Received message without 'req-info', skipping.")
|
||||
in_message.ack(False)
|
||||
self.future.set_result(False)
|
||||
return self.future
|
||||
|
||||
for file_key in req_info['files']:
|
||||
files_info = req_info['files'].get(file_key, []) # Handle missing file types
|
||||
self.files_to_download.extend(files_info)
|
||||
|
||||
if not self.files_to_download:
|
||||
print("No files to download in this message.")
|
||||
self.future.set_result(False)
|
||||
return self.future
|
||||
|
||||
# Create a thread for each file download
|
||||
output_dir = "downloaded_files" # You can change this
|
||||
for file_info in self.files_to_download:
|
||||
thread = threading.Thread(
|
||||
target=self.download_file,
|
||||
args=(in_message.channel, file_info, message_data, output_dir)
|
||||
)
|
||||
self.threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Wait for all download threads to complete
|
||||
asyncio.run(self.wait_for_download_completion(self.threads))
|
||||
|
||||
# After all files are downloaded, acknowledge the original message
|
||||
print("All files are being downloaded. Acknowledging the main message.")
|
||||
in_message.ack(False)
|
||||
print(f"Message from {in_message.routing_key} acknowledged.")
|
||||
return self.future
|
||||
|
||||
async def wait_for_download_completion(self, threads):
|
||||
"""
|
||||
Waits for all download threads to complete.
|
||||
|
||||
Args:
|
||||
threads (list): A list of threads to wait for.
|
||||
"""
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
self.future.set_result(True)
|
||||
@@ -58,7 +58,8 @@ class Dispatch:
|
||||
self.service_port = config.getint('CleverMicro-AMQ', self.PREFIX + 'service-port', fallback=8080)
|
||||
self.rabbit_mq_user = os.getenv("RABBIT_MQ_USER", "springuser")
|
||||
self.rabbit_mq_password = os.getenv("RABBIT_MQ_PASSWORD", "TheCleverWho")
|
||||
|
||||
self.rabbit_mq_url = f"amqp://{self.rabbit_mq_user}:{self.rabbit_mq_password}@{self.amq_host}/"
|
||||
self.download_dir = config.get('CleverMicro-AMQ', self.PREFIX + 'download-dir', fallback='downloaded_files')
|
||||
|
||||
class Backpressure:
|
||||
"""
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# CleverMicro AMQ Adapter settings
|
||||
# ====================================================================
|
||||
[CleverMicro-AMQ]
|
||||
cm.amq-adapter.service-name=amq-adapter
|
||||
cm.amq-adapter.service-name=amq-adapter-python
|
||||
cm.amq-adapter.is-router=false
|
||||
cm.amq-adapter.generator-id=164
|
||||
cm.amq-adapter.route-mapping=amq-cleverswarm::amq-cleverswarm:cleverswarm.localhost.#:5:0
|
||||
|
||||
@@ -2,10 +2,10 @@
|
||||
# CleverMicro AMQ Adapter settings
|
||||
# ====================================================================
|
||||
[CleverMicro-AMQ]
|
||||
cm.amq-adapter.service-name=amq-adapter
|
||||
cm.amq-adapter.service-name=amq-adapter-python
|
||||
cm.amq-adapter.is-router=false
|
||||
cm.amq-adapter.generator-id=211
|
||||
cm.amq-adapter.route-mapping=amq-cleverswarm::amq-cleverswarm:cleverswarm.localhost.#:5:0
|
||||
cm.amq-adapter.route-mapping=amq-cleverswarm::amq-cleverswarm:localhost.#:5:0
|
||||
|
||||
cm.dispatch.use-dlq=true
|
||||
cm.dispatch.use-confirms=false
|
||||
|
||||
+85
-1
@@ -11,10 +11,86 @@ RS = ":"
|
||||
|
||||
|
||||
"""
|
||||
Define the structure of messages used in Clever Microś
|
||||
Define the structure of messages used in Clever Micro
|
||||
"""
|
||||
|
||||
|
||||
class CleverMicroMessage:
|
||||
DEFAULT_CONTENT_TYPE = ["application/octet-stream"]
|
||||
|
||||
def __init__(self, magic: str, id: SnowflakeId, m_field: str, d_field: str, p_field: str,
|
||||
headers: Dict[str, List[str]], trace_info: Dict[str, str], base64body: bytes):
|
||||
self._magic = magic
|
||||
self._id: SnowflakeId = id
|
||||
self._m_field = m_field
|
||||
self._d_field = d_field
|
||||
self._p_field = p_field
|
||||
self._headers: Dict[str, List[str]] = headers
|
||||
self._trace_info: Dict[str, str] = trace_info
|
||||
self._base64body: bytes = base64body
|
||||
|
||||
def magic(self):
|
||||
return self._magic
|
||||
|
||||
def id(self):
|
||||
return self._id
|
||||
|
||||
def m_field(self) -> str:
|
||||
return self._m_field
|
||||
|
||||
def d_field(self) -> str:
|
||||
return self._d_field
|
||||
|
||||
def p_field(self) -> str:
|
||||
return self._p_field
|
||||
|
||||
def headers(self) -> Dict[str, List[str]]:
|
||||
return self._headers
|
||||
|
||||
def trace_info(self) -> Dict[str, str]:
|
||||
return self._trace_info
|
||||
|
||||
def base64body(self) -> bytes:
|
||||
return self._base64body
|
||||
|
||||
def body(self) -> bytes:
|
||||
return base64.b64decode(self.base64body())
|
||||
|
||||
def contentType(self) -> str:
|
||||
return self.headers().get("Content-Type", self.DEFAULT_CONTENT_TYPE)[0]
|
||||
|
||||
|
||||
class DataMessage(CleverMicroMessage):
|
||||
def __init__(self, magic: str, id: SnowflakeId, method: str, domain: str, path: str,
|
||||
headers: Dict[str, List[str]], trace_info: Dict[str, str], base64body: bytes):
|
||||
super().__init__(magic, id, method, domain, path, headers, trace_info, base64body)
|
||||
|
||||
def method(self) -> str:
|
||||
return self.m_field()
|
||||
|
||||
def domain(self) -> str:
|
||||
return self.d_field()
|
||||
|
||||
def path(self) -> str:
|
||||
return self.p_field()
|
||||
|
||||
|
||||
class DataResponse(CleverMicroMessage):
|
||||
def __init__(self, magic: str, id: SnowflakeId, response_code: str, error: str, error_cause: str,
|
||||
headers: Dict[str, List[str]], trace_info: Dict[str, str], base64body: bytes):
|
||||
super().__init__(magic, id, response_code, error, error_cause, headers, trace_info, base64body)
|
||||
|
||||
def response_code(self) -> str:
|
||||
return self.m_field()
|
||||
|
||||
def error(self) -> str:
|
||||
return self.d_field()
|
||||
|
||||
def error_cause(self) -> str:
|
||||
return self.p_field()
|
||||
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AMQRoute:
|
||||
"""
|
||||
@@ -86,6 +162,14 @@ class AMQMessage:
|
||||
return base64.b64decode(self.base64_body)
|
||||
|
||||
|
||||
class MultipartAMQMessage(AMQMessage):
|
||||
def __init__(self, message: AMQMessage, extra_data: dict):
|
||||
super().__init__(message.magic, message.id,
|
||||
message.method, message.domain, message.path,
|
||||
message.headers, message.trace_info,
|
||||
message.base64_body)
|
||||
self.extra_data = extra_data
|
||||
|
||||
@dataclass
|
||||
class AMQResponse:
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
import io
|
||||
|
||||
|
||||
class UploadedFile:
|
||||
def __init__(self, filename: str = None):
|
||||
self.filename = filename
|
||||
self.body_b64 = io.BytesIO()
|
||||
|
||||
|
||||
def from_bytes(input_bytes: bytes) -> UploadedFile:
|
||||
"""
|
||||
Builds the UploadedFile from its serialized (byte array) form.
|
||||
:param input_bytes: serialized message
|
||||
:return: UploadedFile instance
|
||||
"""
|
||||
prolog_len = (input_bytes[0] << 8) + input_bytes[1]
|
||||
metadata = input_bytes[2:prolog_len + 2].decode()
|
||||
body_b64 = input_bytes[prolog_len + 2:]
|
||||
|
||||
return UploadedFile()
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import logging
|
||||
import pika
|
||||
from pika.exchange_type import ExchangeType
|
||||
|
||||
from aio_pika.abc import AbstractRobustQueue, AbstractRobustExchange
|
||||
from pika.exchange_type import ExchangeType
|
||||
from amqp.model.model import AMQRoute
|
||||
from amqp.rabbitmq.rabbit_mq_producer import RabbitMQProducer
|
||||
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE, DLQ_EXCHANGE, AMQ_RPC_EXCHANGE
|
||||
@@ -16,124 +16,66 @@ class RabbitMQConsumer(RabbitMQProducer):
|
||||
def __init__(self, configuration):
|
||||
super().__init__(configuration, RouterConsumer(configuration))
|
||||
|
||||
self.reply_to_exchange: AbstractRobustExchange | None = None # when sending response to RPC request
|
||||
self.reply_to_queue: AbstractRobustQueue | None= None # when receiving response to RPC request
|
||||
|
||||
async def init(self, loop) -> AbstractRobustExchange:
|
||||
await self.setup_amq_channel(loop)
|
||||
await self.router.init_routing_paths_consumer(self.router.channel)
|
||||
cm_reply_to_queue = self.router.get_reply_to_queue_name()
|
||||
self.reply_to_queue = await self.channel.declare_queue(
|
||||
name=cm_reply_to_queue, durable=True, exclusive=False, auto_delete=False
|
||||
)
|
||||
self.reply_to_exchange = await self.channel.declare_exchange(
|
||||
AMQ_REPLY_TO_EXCHANGE, type=ExchangeType.direct
|
||||
)
|
||||
await self.reply_to_queue.bind(exchange=self.reply_to_exchange, routing_key=cm_reply_to_queue)
|
||||
return self.reply_to_exchange
|
||||
|
||||
def cleanup(self):
|
||||
logging.info(" [*] RabbitMQConsumer.cleanup()")
|
||||
self.router.cleanup() # de-register route(s)
|
||||
if self.router.is_open(self.channel):
|
||||
try:
|
||||
self.channel.close()
|
||||
except (pika.connection.exceptions.AMQPError, pika.connection.exceptions.ChannelClosedByBroker) as e:
|
||||
except Exception as e:
|
||||
logging.error(" [E] RabbitMQConsumer.PreDestroy cleanup - connection close error: %s", e)
|
||||
|
||||
def register_inbound_routes(self, callback):
|
||||
async def register_inbound_routes(self, route_mapping: str, callback):
|
||||
if self.router.is_open(self.channel):
|
||||
queue_args = self.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
|
||||
route_mapping = self.amq_configuration.amq_adapter.route_mapping
|
||||
logging.info(" [*] RabbitMQConsumer register routes, cfg mapping: [%s]", route_mapping)
|
||||
requested_routes = self.router.unwrap_route_list(route_mapping)
|
||||
for route in requested_routes:
|
||||
try:
|
||||
queue_name = route.queue if route.queue else self.router.get_consume_queue_name()
|
||||
logging.info(" [*] RabbitMQConsumer register route: [%s]", route)
|
||||
self.channel.queue_declare(queue_name, durable=True, exclusive=False,
|
||||
auto_delete=False, arguments=queue_args)
|
||||
exchange = route.exchange if route.exchange else AMQ_RPC_EXCHANGE
|
||||
self.channel.exchange_declare(exchange, exchange_type=pika.exchange_type.ExchangeType.topic)
|
||||
self.channel.queue_bind(queue_name, exchange, route.key)
|
||||
self.channel.basic_consume(
|
||||
queue=queue_name,
|
||||
on_message_callback=callback,
|
||||
auto_ack=False
|
||||
queue: AbstractRobustQueue = await self.channel.declare_queue(name=queue_name,
|
||||
durable=True,
|
||||
exclusive=False,
|
||||
auto_delete=False,
|
||||
arguments=queue_args)
|
||||
exchange_name: str = route.exchange if route.exchange else AMQ_RPC_EXCHANGE
|
||||
exchange: AbstractRobustExchange = await self.channel.declare_exchange(
|
||||
name=exchange_name, type=ExchangeType.topic
|
||||
)
|
||||
self.router.register_route(
|
||||
await queue.bind(exchange, route.key)
|
||||
await queue.consume(callback, no_ack=False)
|
||||
await self.router.register_route(
|
||||
AMQRoute(
|
||||
route.component_name, exchange, queue_name,
|
||||
route.component_name, exchange_name, queue_name,
|
||||
sanitize_as_rabbitmq_name(route.key),
|
||||
route.timeout, route.valid_until
|
||||
)
|
||||
)
|
||||
logging.info(" [*] AMQ connection initialized for route: %s, listens on: %s/%s",
|
||||
route, exchange, queue_name)
|
||||
except pika.connection.exceptions.AMQPError as err:
|
||||
route, exchange, queue_name)
|
||||
except Exception as err:
|
||||
logging.error(" [E] Callback registration INCOMPLETE, %s", err)
|
||||
else:
|
||||
logging.warning(" [W] Channel is null, cannot register routes.")
|
||||
|
||||
def register_rpc_handler(self, rpc_consumer):
|
||||
cm_reply_to_queue = self.router.get_reply_to_queue_name()
|
||||
self.channel.queue_declare(cm_reply_to_queue, durable=True, exclusive=False, auto_delete=False)
|
||||
self.channel.exchange_declare(
|
||||
AMQ_REPLY_TO_EXCHANGE, exchange_type=pika.exchange_type.ExchangeType.direct
|
||||
async def register_rpc_handler(self, rpc_consumer):
|
||||
await self.reply_to_queue.consume(callback=rpc_consumer,
|
||||
no_ack=False
|
||||
)
|
||||
self.channel.queue_bind(cm_reply_to_queue, AMQ_REPLY_TO_EXCHANGE, cm_reply_to_queue)
|
||||
self.channel.basic_consume(
|
||||
queue=cm_reply_to_queue,
|
||||
on_message_callback=rpc_consumer,
|
||||
auto_ack=False
|
||||
)
|
||||
|
||||
"""
|
||||
def limit_rate_inside_time_window(self, message: AMQMessage) -> int:
|
||||
next_batch_size = BackpressureSubscriber.NO_OP
|
||||
now = time.time_ns() // 1000000
|
||||
queue_size = self.getQueueSize(message)
|
||||
|
||||
# 1. Backpressure? if unprocessed count exceeds threshold, wait for some time
|
||||
if queue_size > THRESHOLD:
|
||||
# TODO - signal queue backpressure explicitly to upstream/metrics
|
||||
# wait: estimated-per-message-processing-time * SQRT(over-the-limit)
|
||||
self.wait_for_milliseconds(
|
||||
(WINDOW / self.request_for_window_duration) * (queue_size - THRESHOLD) ** 0.5)
|
||||
queue_size = self.getQueueSize(message) # refresh queue size after wait
|
||||
|
||||
# 2. batch complete? determine next batch size and wait till window end. reset window.
|
||||
if self.fulfilled >= self.request_for_window_duration:
|
||||
next_batch_size = self.calculateNextBatchSize(queue_size)
|
||||
self.request_for_window_duration = next_batch_size
|
||||
remaining_time = (self.window_started_at + WINDOW) - now
|
||||
if remaining_time > 0:
|
||||
self.wait_for_milliseconds(remaining_time)
|
||||
queue_size = self.getQueueSize(message) # refresh queue size after wait
|
||||
self.window_started_at = time.time_ns() // 1000000
|
||||
self.fulfilled = 0
|
||||
|
||||
# 3. Window elapsed? Reset window & recalculate batch size
|
||||
if now >= self.window_started_at + WINDOW:
|
||||
if next_batch_size == BackpressureSubscriber.NO_OP:
|
||||
next_batch_size = self.calculateNextBatchSize(queue_size)
|
||||
self.request_for_window_duration = next_batch_size
|
||||
self.window_started_at = time.time_ns() // 1000000
|
||||
self.fulfilled = 0
|
||||
|
||||
# 4. send the message from the current batch
|
||||
self.send_message(message)
|
||||
self.fulfilled += 1
|
||||
|
||||
# 5. report and setup state for the next call
|
||||
end = time.time_ns() // 1000000
|
||||
self.on_next_timestamp = end
|
||||
return next_batch_size
|
||||
|
||||
def calculateNextBatchSize(self, queue_size: int) -> int:
|
||||
next_batch_size = self.request_for_window_duration
|
||||
if queue_size <= THRESHOLD:
|
||||
# processing proceeds well, can increase the rate
|
||||
next_batch_size += self.subscriber.increaseRate(self.request_for_window_duration)
|
||||
logging.debug(f"INCREASE rate to {next_batch_size}. qSize={queue_size}")
|
||||
else:
|
||||
decrease = self.subscriber.backoff(queue_size - THRESHOLD)
|
||||
if self.request_for_window_duration > decrease:
|
||||
next_batch_size -= decrease
|
||||
logging.debug(f"DECREASE rate to {next_batch_size}. qSize={queue_size}")
|
||||
return next_batch_size
|
||||
"""
|
||||
"""
|
||||
def getQueueSize(self, message: AMQMessage) -> int:
|
||||
try:
|
||||
route_queue_name = self.router.findRoute(message)
|
||||
queue_name = route_queue_name.queue if route_queue_name else self.TASK_DISPATCH_QUEUE_NAME
|
||||
logging.info(f"Get Queue size for {queue_name}")
|
||||
return max(len(self.awaitingACK), self.channel.message_count(queue_name))
|
||||
except Exception:
|
||||
return len(self.awaitingACK)
|
||||
"""
|
||||
|
||||
@@ -1,10 +1,5 @@
|
||||
import logging
|
||||
import time
|
||||
|
||||
import pika
|
||||
from pika.adapters.blocking_connection import BlockingChannel, BlockingConnection
|
||||
from pika.channel import Channel
|
||||
from pika.connection import Connection
|
||||
import aio_pika
|
||||
|
||||
from amqp.adapter.amq_message_factory import AMQMessageFactory
|
||||
from amqp.config.amq_configuration import AMQConfiguration
|
||||
@@ -14,38 +9,22 @@ from amqp.router.router_producer import RouterProducer
|
||||
|
||||
class RabbitMQProducer:
|
||||
PREFETCH_COUNT = 1
|
||||
THRESHOLD = 1000
|
||||
WINDOW = 5000 # 5 seconds
|
||||
|
||||
def __init__(self, amq_configuration: AMQConfiguration, router: RouterProducer = None):
|
||||
self.channel: BlockingChannel | None = None
|
||||
self.connection: BlockingConnection | None = None
|
||||
self.request_for_window_duration = None
|
||||
self.channel: aio_pika.RobustChannel | None = None
|
||||
self.connection: aio_pika.RobustConnection | None = None
|
||||
self.rpc_exchange: aio_pika.RobustExchange | None = None
|
||||
self.reply_exchange: aio_pika.RobustExchange | None = None
|
||||
self.amq_configuration = amq_configuration
|
||||
logging.info("*******************************************")
|
||||
logging.info("******** RabbitMQProducer.init() ********** %s - %s - %s",
|
||||
self.amq_configuration.dispatch.amq_host,
|
||||
self.amq_configuration.dispatch.use_confirms,
|
||||
self.WINDOW)
|
||||
logging.info("******** RabbitMQProducer.init(): %s",
|
||||
self.amq_configuration.dispatch.amq_host)
|
||||
logging.info("*******************************************")
|
||||
self.router = router or RouterProducer(amq_configuration)
|
||||
|
||||
# this.meterRegistry = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT);
|
||||
|
||||
self.setup_amq_channel()
|
||||
self.router.init_routing_paths(self.channel)
|
||||
|
||||
# self.backpressure_per_exchange = defaultdict(lambda: Gauge(
|
||||
# f"amqp.adapter.exchange.{exchange}", registry=CollectorRegistry()))
|
||||
|
||||
self.router.set_route_added_notifier(
|
||||
lambda route: logging.info(" [DBG] ---------- RouteAdded, setting up metrics: amqp.adapter.exchange.%s",
|
||||
route.exchange))
|
||||
|
||||
self.request_for_window_duration = self.THRESHOLD // 2
|
||||
self.window_started_at = time.time()
|
||||
self.on_next_timestamp = self.window_started_at
|
||||
self.fulfilled = 0
|
||||
async def init(self, loop):
|
||||
await self.setup_amq_channel(loop)
|
||||
|
||||
def cleanup(self):
|
||||
if self.connection:
|
||||
@@ -54,30 +33,27 @@ class RabbitMQProducer:
|
||||
except Exception as e:
|
||||
logging.error("RabbitMQConsumer.PreDestroy cleanup - connection close error: %s", e)
|
||||
|
||||
def setup_amq_channel(self):
|
||||
async def setup_amq_channel(self, loop):
|
||||
"""
|
||||
Set up the RabbitMQ connection
|
||||
:return: nothing, it applies changes to instance variables
|
||||
"""
|
||||
credentials: pika.credentials.PlainCredentials = pika.credentials.PlainCredentials(
|
||||
username=self.amq_configuration.dispatch.rabbit_mq_user,
|
||||
password=self.amq_configuration.dispatch.rabbit_mq_password
|
||||
)
|
||||
factory = pika.ConnectionParameters(
|
||||
self.connection = await aio_pika.connect_robust(
|
||||
host=self.amq_configuration.dispatch.amq_host,
|
||||
credentials=credentials
|
||||
port=self.amq_configuration.dispatch.amq_port,
|
||||
login=self.amq_configuration.dispatch.rabbit_mq_user,
|
||||
password=self.amq_configuration.dispatch.rabbit_mq_password,
|
||||
loop=loop
|
||||
)
|
||||
self.connection = pika.BlockingConnection(factory)
|
||||
logging.info("RabbitMQProducer.setupAMQ, got connection, open=%s", self.connection.is_open)
|
||||
self.connection.add_on_connection_blocked_callback(self.blocked_listener)
|
||||
self.connection.add_on_connection_unblocked_callback(self.unblocked_listener)
|
||||
self.channel = self.connection.channel()
|
||||
self.channel.basic_qos(prefetch_count=self.PREFETCH_COUNT)
|
||||
if self.amq_configuration.dispatch.use_confirms:
|
||||
self.setup_local_ack()
|
||||
self.channel.add_on_return_callback(self.return_listener)
|
||||
logging.info("RabbitMQProducer.setupAMQ, got connection, open=%s", self.connection.is_closed)
|
||||
self.channel = await self.connection.channel()
|
||||
await self.channel.set_qos(prefetch_count=self.PREFETCH_COUNT)
|
||||
await self.router.init_routing_paths(self.channel)
|
||||
self.router.set_route_added_notifier(
|
||||
lambda route: logging.info(" [DBG] ---------- RouteAdded, setting up metrics: amqp.adapter.exchange.%s",
|
||||
route.exchange))
|
||||
|
||||
def send_message(self, message: AMQMessage):
|
||||
async def send_message(self, message: AMQMessage):
|
||||
"""
|
||||
The API routine that sends a datagram message (AMQMessage) to the destination. The Destination
|
||||
is determined from the AMQMessage metadata section, using `domain` and `path` as key.
|
||||
@@ -87,27 +63,31 @@ class RabbitMQProducer:
|
||||
"""
|
||||
try:
|
||||
route = self.router.find_route_by_message(message)
|
||||
if route:
|
||||
if route.timeout < 0:
|
||||
if route and self.rpc_exchange:
|
||||
if route.timeout <= 0:
|
||||
# No response expected
|
||||
self.channel.basic_publish(
|
||||
route.exchange,
|
||||
self.router.get_routing_key(message), # context path is a routing key
|
||||
properties=pika.BasicProperties(delivery_mode=2),
|
||||
body=AMQMessageFactory.serialize(message)
|
||||
pika_message: aio_pika.message.Message = aio_pika.message.Message(
|
||||
body=AMQMessageFactory.serialize(message),
|
||||
content_encoding='utf-8',
|
||||
delivery_mode=2
|
||||
)
|
||||
await self.rpc_exchange.publish(
|
||||
message=pika_message,
|
||||
routing_key=self.router.get_routing_key(message), # context path is a routing key
|
||||
)
|
||||
logging.info(f" [x] basicPublish (no-reply), messageId: {message.id}, to: {route}")
|
||||
else:
|
||||
# This is RPC call, there should be response
|
||||
props = pika.BasicProperties(
|
||||
correlation_id=message.id,
|
||||
pika_message: aio_pika.message.Message = aio_pika.message.Message(
|
||||
body=AMQMessageFactory.serialize(message),
|
||||
content_encoding='utf-8',
|
||||
delivery_mode=2,
|
||||
correlation_id=str(message.id),
|
||||
reply_to=self.router.get_reply_to_queue_name()
|
||||
)
|
||||
self.channel.basic_publish(
|
||||
route.exchange,
|
||||
self.router.get_routing_key(message), # context path is a routing key
|
||||
properties=props,
|
||||
body=AMQMessageFactory.serialize(message)
|
||||
await self.rpc_exchange.publish(
|
||||
message=pika_message,
|
||||
routing_key=self.router.get_routing_key(message), # context path is a routing key
|
||||
)
|
||||
logging.info(" [x] basicPublish (RPC call), messageId: %s, to: %s",
|
||||
message.id.as_string(), route.as_string())
|
||||
@@ -117,30 +97,8 @@ class RabbitMQProducer:
|
||||
except Exception as err:
|
||||
logging.error("Send message failed: %s", err)
|
||||
|
||||
def request_routes(self):
|
||||
self.router.request_routes_from_adapters()
|
||||
|
||||
def setup_local_ack(self):
|
||||
self.channel.confirm_delivery()
|
||||
async def request_routes(self):
|
||||
await self.router.request_routes_from_adapters()
|
||||
|
||||
def find_route(self, domain: str, path: str) -> AMQRoute:
|
||||
return self.router.find_route(domain, path)
|
||||
|
||||
def return_listener(self, return_message):
|
||||
logging.error(
|
||||
f" [E] ****** (NOT DELIVERED) ID={return_message.properties.message_id}, "
|
||||
f"Exchg=[{return_message.exchange}], Key=[{return_message.routing_key}], "
|
||||
f"replyTxt={return_message.reply_text}, length={len(return_message.body) if return_message.body else 0}"
|
||||
)
|
||||
|
||||
def blocked_listener(self, reason: str):
|
||||
logging.error(f" [E] CHANNEL BLOCKED, reason={reason}")
|
||||
|
||||
def unblocked_listener(self):
|
||||
logging.info(" [*] CHANNEL UNBLOCKED")
|
||||
|
||||
def ack_nack_callback(method):
|
||||
if method == pika.spec.Basic.Ack:
|
||||
pass
|
||||
if method == pika.spec.Basic.Nack:
|
||||
pass
|
||||
@@ -1,9 +1,7 @@
|
||||
import logging
|
||||
from typing import Set
|
||||
|
||||
from pika import spec
|
||||
from pika.channel import Channel
|
||||
|
||||
from aio_pika.abc import AbstractRobustChannel, AbstractRobustQueue, AbstractIncomingMessage
|
||||
from amqp.adapter.amq_route_factory import AMQRouteFactory
|
||||
from amqp.config.amq_configuration import AMQConfiguration
|
||||
from amqp.model.model import CleverMicroServiceMessage, AMQRoute
|
||||
@@ -30,25 +28,21 @@ class RouterConsumer(RouterProducer):
|
||||
self.amqRouteFactory = AMQRouteFactory()
|
||||
logging.info(" [*] RouterConsumer.<init>")
|
||||
|
||||
def init_routing_paths_consumer(self, _channel: Channel):
|
||||
async def init_routing_paths_consumer(self, _channel: AbstractRobustChannel):
|
||||
"""
|
||||
* Initialize router according to Consumer mode.
|
||||
"""
|
||||
|
||||
self.init_routing_paths_consumer(_channel)
|
||||
logging.info(" [*] RouterConsumer.initConsumerRouter, channel open: %s", self.is_open(_channel))
|
||||
if self.is_open(_channel):
|
||||
try:
|
||||
# 1. declare a Queue that gives the RoutingDAta requests
|
||||
cm_request_queue: str = self.amq_configuration.amq_adapter.adapter_prefix() + CM_C_REQ_QUEUE
|
||||
self.channel.queue_declare(cm_request_queue, durable=True, exclusive=False, auto_delete=False)
|
||||
queue: AbstractRobustQueue = await self.channel.declare_queue(name=cm_request_queue, durable=True, exclusive=False, auto_delete=False)
|
||||
# 1a. bind queue to CM_RESPONSE_EXCHANGE w/o any conditions (empty routing key)
|
||||
self.channel.queue_bind(queue=cm_request_queue, exchange=CM_REQUEST_EXCHANGE, routing_key="")
|
||||
await queue.bind(exchange=CM_REQUEST_EXCHANGE, routing_key="")
|
||||
# 1b. Service adapter mode - listens on Request queue and replies to Response queue
|
||||
logging.info(" [*] RouterConsumer listens for RouteDataRequests on: %s", cm_request_queue)
|
||||
self.channel.basic_consume(queue=cm_request_queue,
|
||||
auto_ack=False,
|
||||
on_message_callback=self.handle_routing_data_request_message)
|
||||
await queue.consume(callback=self.handle_routing_data_request_message, no_ack=False)
|
||||
except Exception as ioe:
|
||||
logging.error(" [E] RouterConsumer FATAL: cannot initialize, %s", ioe)
|
||||
else:
|
||||
@@ -58,39 +52,35 @@ class RouterConsumer(RouterProducer):
|
||||
* ServiceMessageType.ROUTING_DATA_REQUEST handler listening on queue 'CleverMicroRequest'
|
||||
"""
|
||||
|
||||
def handle_routing_data_request_message(self,
|
||||
cb_channel: Channel,
|
||||
method: spec.Basic.Deliver,
|
||||
properties: spec.BasicProperties,
|
||||
body: bytes):
|
||||
async def handle_routing_data_request_message(self, message: AbstractIncomingMessage):
|
||||
# some debug statements, TODO - remove for prod release
|
||||
delivery_tag = method.delivery_tag
|
||||
debug = "Delivery.Properties = " + properties.__str__()
|
||||
logging.info(" [DBG] ******[%] (ROUTING_DATA_REQ) tag: %s, env: %s, props: %s",
|
||||
delivery_tag, method.consumer_tag, method.get_properties(), debug)
|
||||
delivery_tag = message.delivery_tag
|
||||
debug = "Delivery.Properties = " + message.properties.__str__()
|
||||
logging.info(" [DBG] ******[%s] (ROUTING_DATA_REQ) tag: %s, env: %s, props: %s",
|
||||
delivery_tag, message.consumer_tag, message.properties, debug)
|
||||
|
||||
cb_channel.basic_ack(delivery_tag, False)
|
||||
await message.ack(False)
|
||||
|
||||
message: CleverMicroServiceMessage = self.service_message_factory.from_bytes(body)
|
||||
cm_message: CleverMicroServiceMessage = self.service_message_factory.from_bytes(message.body)
|
||||
route_mapping: str = self.wrap_own_routes()
|
||||
logging.info(" [DBG] ******[%s] (ROUTING_DATA_REQ) msg=%s, routeMap=[%s]",
|
||||
delivery_tag, self.service_message_factory.to_string(message), route_mapping)
|
||||
delivery_tag, self.service_message_factory.to_string(cm_message), route_mapping)
|
||||
# Reply only if self adapter has own routeMapping (e.g. API-gateway-side adapter has none)
|
||||
if message.is_valid() and message.message_type == ServiceMessageType.ROUTING_DATA_REQUEST and route_mapping:
|
||||
if cm_message.is_valid() and cm_message.message_type == ServiceMessageType.ROUTING_DATA_REQUEST and route_mapping:
|
||||
try:
|
||||
service_message: CleverMicroServiceMessage = self.service_message_factory.of(
|
||||
ServiceMessageType.ROUTING_DATA, route_mapping, message.reply_to
|
||||
ServiceMessageType.ROUTING_DATA, route_mapping, cm_message.reply_to
|
||||
)
|
||||
self.publish_service_message(service_message, CM_RESPONSE_EXCHANGE)
|
||||
await self.publish_service_message(service_message, self.cm_response_exchange)
|
||||
except Exception as err:
|
||||
logging.error(" [E] ******[%s] (ROUTING_DATA_REQ) Response to queue: %s failed: %s",
|
||||
delivery_tag, CM_RESPONSE_EXCHANGE, err)
|
||||
|
||||
else:
|
||||
logging.error(" [E] ******[%s] (ROUTING_DATA_REQ) [%s] produces invalid or no CM msg type.",
|
||||
delivery_tag, str(body))
|
||||
delivery_tag, str(message.body))
|
||||
|
||||
def register_route(self, route: AMQRoute):
|
||||
async def register_route(self, route: AMQRoute):
|
||||
"""
|
||||
* Called in Consumer (Service Adapter) context. Registers the routing data for self Service
|
||||
* with the Master router by sending own routing data to Master via CleverMicroResponse queue.
|
||||
@@ -101,7 +91,7 @@ class RouterConsumer(RouterProducer):
|
||||
serviceMessage: CleverMicroServiceMessage = self.service_message_factory.of(
|
||||
ServiceMessageType.ROUTING_DATA, route.as_string(), ANY_RECIPIENT
|
||||
)
|
||||
self.publish_service_message(serviceMessage, CM_RESPONSE_EXCHANGE)
|
||||
await self.publish_service_message(serviceMessage, self.cm_response_exchange)
|
||||
self.add_own_route(route)
|
||||
logging.info(" [DBG] RouterConsumer registered Route: %s", route)
|
||||
else:
|
||||
@@ -110,7 +100,7 @@ class RouterConsumer(RouterProducer):
|
||||
except Exception as ioe:
|
||||
logging.error(" [E] RouterConsumer failed register route data: %s", ioe)
|
||||
|
||||
def remove_route_on_cleanup(self, route: AMQRoute) -> bool:
|
||||
async def remove_route_on_cleanup(self, route: AMQRoute) -> bool:
|
||||
"""
|
||||
* Unregister / remove route from the list
|
||||
* :param route: route to remove
|
||||
@@ -123,7 +113,7 @@ class RouterConsumer(RouterProducer):
|
||||
route.as_string(),
|
||||
ANY_RECIPIENT
|
||||
)
|
||||
if not self.publish_service_message(service_message, CM_RESPONSE_EXCHANGE):
|
||||
if not await self.publish_service_message(service_message, self.cm_response_exchange):
|
||||
logging.warning(" [W] RouterConsumer couldn't remove current route: %s, channel not open.",
|
||||
route)
|
||||
|
||||
|
||||
@@ -7,10 +7,10 @@
|
||||
* forwards it to the application.
|
||||
"""
|
||||
import logging
|
||||
import os
|
||||
import pika
|
||||
from pika import BasicProperties, spec
|
||||
from pika.channel import Channel
|
||||
|
||||
from aio_pika import Message
|
||||
from aio_pika.abc import AbstractRobustChannel, AbstractRobustQueue, AbstractRobustConnection, AbstractRobustExchange, \
|
||||
AbstractIncomingMessage
|
||||
from pika.exchange_type import ExchangeType
|
||||
|
||||
from amqp.adapter.service_message_factory import CleverMicroServiceMessageFactory
|
||||
@@ -29,26 +29,18 @@ class RouterProducer(RouterBase):
|
||||
* :param configuration: adapter's configuration reference.
|
||||
"""
|
||||
super().__init__()
|
||||
self.connection = None
|
||||
self.connection: AbstractRobustConnection | None = None
|
||||
self.service_message_factory = CleverMicroServiceMessageFactory(configuration.amq_adapter.service_name)
|
||||
self.amq_configuration: AMQConfiguration = configuration
|
||||
self.channel: Channel | None = None
|
||||
self.channel: AbstractRobustChannel | None = None
|
||||
self.cm_request_exchange: AbstractRobustExchange | None = None
|
||||
self.cm_response_exchange: AbstractRobustExchange | None = None
|
||||
|
||||
logging.info(f" [*] RouterProducer.<init>: %s, route: %s",
|
||||
self.amq_configuration.amq_adapter.service_name,
|
||||
self.amq_configuration.amq_adapter.route_mapping)
|
||||
|
||||
def amq_consumer(self):
|
||||
credentials = pika.PlainCredentials(
|
||||
os.environ['RABBIT_MQ_USER'], os.environ['RABBIT_MQ_PASSWORD']
|
||||
)
|
||||
self.connection: pika.SelectConnection = pika.SelectConnection(
|
||||
pika.ConnectionParameters(host=self.amq_configuration.dispatch.amq_host, credentials=credentials)
|
||||
)
|
||||
self.channel = self.connection.channel()
|
||||
# self.channel.basic_consume()
|
||||
|
||||
def init_routing_paths(self, channel: Channel | None):
|
||||
async def init_routing_paths(self, channel: AbstractRobustChannel | None):
|
||||
"""
|
||||
* Initialize router according to Producer mode.
|
||||
* The logic here sets a producer and listener on CleverMicroRequest/CleverMicroResponse queues
|
||||
@@ -63,22 +55,25 @@ class RouterProducer(RouterBase):
|
||||
try:
|
||||
# **** set up the service discovery internal exchanges ****
|
||||
# 1. declare RequestExchange to where to publish the RoutingData requests
|
||||
self.channel.exchange_declare(exchange=CM_REQUEST_EXCHANGE, exchange_type=ExchangeType.fanout)
|
||||
self.cm_request_exchange = await self.channel.declare_exchange(
|
||||
name=CM_REQUEST_EXCHANGE, type=ExchangeType.fanout
|
||||
)
|
||||
#
|
||||
# 2. declare ResponseExchange where to publish current RoutingData (response to
|
||||
# a RoutingData request)
|
||||
self.channel.exchange_declare(exchange=CM_RESPONSE_EXCHANGE, exchange_type=ExchangeType.fanout)
|
||||
self.cm_response_exchange = await self.channel.declare_exchange(
|
||||
name=CM_RESPONSE_EXCHANGE, type=ExchangeType.fanout
|
||||
)
|
||||
# 2a. declare a Response queue for Routing Data replies from Services
|
||||
cm_response_queue: str = self.get_response_queue_name()
|
||||
self.channel.queue_declare(cm_response_queue, durable=True, exclusive=False, auto_delete=False,
|
||||
arguments={})
|
||||
response_queue: AbstractRobustQueue = await self.channel.declare_queue(
|
||||
name=cm_response_queue, durable=True, exclusive=False, auto_delete=False,
|
||||
arguments={}
|
||||
)
|
||||
# 2b. bind queue to CM_RESPONSE_EXCHANGE w/o any conditions (empty routing key)
|
||||
self.channel.queue_bind(queue=cm_response_queue, exchange=CM_RESPONSE_EXCHANGE, routing_key="#")
|
||||
await response_queue.bind(exchange=CM_RESPONSE_EXCHANGE, routing_key="#")
|
||||
# 2c. listen for incoming RoutingData messages
|
||||
self.channel.basic_consume(
|
||||
queue=cm_response_queue,
|
||||
auto_ack=False,
|
||||
on_message_callback=self.handle_routing_data_message)
|
||||
await response_queue.consume(callback=self.handle_routing_data_message, no_ack=False)
|
||||
logging.info(" [*] Routing master listens for Route updates on %s", cm_response_queue)
|
||||
except Exception as ioe:
|
||||
logging.error(" [E] RouterProducer FATAL: cannot initialize, %s", ioe)
|
||||
@@ -86,38 +81,34 @@ class RouterProducer(RouterBase):
|
||||
else:
|
||||
logging.error(" [E] RouterProducer FATAL: cannot initialize, RabbitMQ channel is not open.")
|
||||
|
||||
def handle_routing_data_message(self,
|
||||
cb_channel: Channel,
|
||||
method: spec.Basic.Deliver,
|
||||
properties: spec.BasicProperties,
|
||||
body: bytes):
|
||||
async def handle_routing_data_message(self, message: AbstractIncomingMessage):
|
||||
"""
|
||||
* Handler for CLEVER_MICRO_RESPONSE queue messages. Used in Master mode to handle receipt of
|
||||
* routing data, invokes addRoute or removeRoute depending on the type of the message.
|
||||
"""
|
||||
logging.info(f" [x] Received {body}, m={method}, p={properties}")
|
||||
logging.info(f" [x] Received {message.body}, m={message.message_id}, p={message.properties}")
|
||||
"""
|
||||
logging.info(" [DBG] ******[%s] ROUTING_DATA %s >>> %s",
|
||||
delivery.getEnvelope().getDeliveryTag(),
|
||||
delivery.getEnvelope().toString(), str(delivery.getBody()))
|
||||
"""
|
||||
# ACK message now so that it is not being redelivered
|
||||
cb_channel.basic_ack(method.delivery_tag, False)
|
||||
message: CleverMicroServiceMessage = self.service_message_factory.from_bytes(body)
|
||||
await message.ack(multiple=False)
|
||||
cm_message: CleverMicroServiceMessage = self.service_message_factory.from_bytes(message.body)
|
||||
# ignore the message if it comes from ourselves (even if from different instance)
|
||||
logging.info(" [DBG] ******[%s] ROUTING_DATA, msg=%s, routeMap=[%s] sn=%s, replyTo=%s",
|
||||
method.delivery_tag,
|
||||
self.service_message_factory.to_string(message),
|
||||
message.delivery_tag,
|
||||
self.service_message_factory.to_string(cm_message),
|
||||
self.amq_configuration.amq_adapter.route_mapping,
|
||||
self.amq_configuration.amq_adapter.service_name,
|
||||
message.reply_to)
|
||||
if message.is_valid() and not self.amq_configuration.amq_adapter.service_name == message.reply_to:
|
||||
if message.message_type == ServiceMessageType.ROUTING_DATA:
|
||||
self.add_routes(message.message)
|
||||
if message.message_type == ServiceMessageType.ROUTING_DATA_REMOVE:
|
||||
self.remove_routes(message.message)
|
||||
cm_message.reply_to)
|
||||
if cm_message.is_valid() and not self.amq_configuration.amq_adapter.service_name == cm_message.reply_to:
|
||||
if cm_message.message_type == ServiceMessageType.ROUTING_DATA:
|
||||
self.add_routes(cm_message.message)
|
||||
if cm_message.message_type == ServiceMessageType.ROUTING_DATA_REMOVE:
|
||||
self.remove_routes(cm_message.message)
|
||||
|
||||
def request_routes_from_adapters(self):
|
||||
async def request_routes_from_adapters(self):
|
||||
"""
|
||||
* Send a broadcast message prompting all listening adapters to reply with own routing data.
|
||||
"""
|
||||
@@ -126,37 +117,37 @@ class RouterProducer(RouterBase):
|
||||
serviceMessage: CleverMicroServiceMessage = self.service_message_factory.of(
|
||||
ServiceMessageType.ROUTING_DATA_REQUEST, "", ANY_RECIPIENT
|
||||
)
|
||||
if not self.publish_service_message(serviceMessage, CM_REQUEST_EXCHANGE):
|
||||
if not await self.publish_service_message(serviceMessage, self.cm_request_exchange):
|
||||
logging.warning(" [E] RouterProducer could not request route data: %s channel is not open.",
|
||||
CM_REQUEST_EXCHANGE)
|
||||
|
||||
except Exception as ioe:
|
||||
logging.error(" [E] RouterProducer failed request routing data: %s", ioe)
|
||||
|
||||
BASIC: BasicProperties = BasicProperties(content_type="application/octet-stream",
|
||||
content_encoding=None,
|
||||
headers=None,
|
||||
delivery_mode=1,
|
||||
priority=0,
|
||||
correlation_id=None)
|
||||
|
||||
def publish_service_message(self, message: CleverMicroServiceMessage, exchange_name: str) -> bool:
|
||||
async def publish_service_message(self,
|
||||
message: CleverMicroServiceMessage,
|
||||
exchange: AbstractRobustExchange) -> bool:
|
||||
"""
|
||||
* Publishes a service message to the service queue and logs success log entry.
|
||||
*
|
||||
* :param message: Service Message to be sent
|
||||
* :param exchangeName: target exchange name where to publish the message
|
||||
* :param queueName target: queue name (Optional, purpose is to get the queue size only)
|
||||
* :param exchange: target exchange name where to publish the message
|
||||
* :return True if channel is open, False otherwise
|
||||
* @throws IOException if something else goes wrong
|
||||
"""
|
||||
if self.is_open(self.channel):
|
||||
binary_content: bytes = self.service_message_factory.serialize(message)
|
||||
self.channel.basic_publish(exchange=exchange_name,
|
||||
routing_key="", # empty routing key
|
||||
properties=self.BASIC,
|
||||
body=binary_content)
|
||||
logging.info(" [x] Service Message Sent to %s, msg: %s", exchange_name, str(binary_content))
|
||||
pika_message: Message = Message(
|
||||
body=binary_content,
|
||||
content_encoding='utf-8',
|
||||
delivery_mode=2,
|
||||
content_type="application/octet-stream",
|
||||
headers=None,
|
||||
priority=0,
|
||||
correlation_id=None
|
||||
)
|
||||
await exchange.publish(message=pika_message, routing_key="")
|
||||
logging.info(" [x] Service Message Sent to %s, msg: %s", exchange.name, str(binary_content))
|
||||
return True
|
||||
return False
|
||||
|
||||
@@ -181,10 +172,10 @@ class RouterProducer(RouterBase):
|
||||
"""
|
||||
return self.amq_configuration.amq_adapter.adapter_prefix() + CM_C_AMQ_QUEUE
|
||||
|
||||
def is_open(self, _channel: Channel) -> bool:
|
||||
def is_open(self, _channel: AbstractRobustChannel) -> bool:
|
||||
"""
|
||||
* Test if the channel is usable (opened).
|
||||
* :param _channel: the channel
|
||||
* :return True if the channel is opened.
|
||||
"""
|
||||
return _channel is not None and _channel.is_open
|
||||
return _channel is not None and not _channel.is_closed
|
||||
|
||||
+33
-12
@@ -1,8 +1,10 @@
|
||||
import logging
|
||||
import logging.config
|
||||
import os
|
||||
import asyncio
|
||||
from asyncio import Future
|
||||
|
||||
from aio_pika.abc import AbstractRobustExchange
|
||||
|
||||
from amqp.adapter.amq_message_factory import AMQMessageFactory
|
||||
from amqp.adapter.amq_route_factory import AMQRouteFactory
|
||||
@@ -29,53 +31,72 @@ class AMQService:
|
||||
CLEVER_AMQP_CLIENT_VERSION = os.getenv("CLEVER_AMQP_CLIENT_VERSION")
|
||||
MAX_WAIT = 300 # defaults to 5 minutes (300 seconds)
|
||||
|
||||
def __init__(self, amq_configuration: AMQConfiguration, http_adapter):
|
||||
def __init__(self, amq_configuration: AMQConfiguration, service_adapter):
|
||||
logging.info("***********************************************")
|
||||
logging.info("********** AMQ Service *******************")
|
||||
logging.info("***********************************************")
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
self.amq_configuration = amq_configuration
|
||||
self.http_adapter = http_adapter
|
||||
self.message_factory = AMQMessageFactory.get_instance(self.amq_configuration.amq_adapter.generator_id)
|
||||
self.rabbit_mq_consumer = RabbitMQConsumer(self.amq_configuration)
|
||||
self.tracer = initialize_trace(name=amq_configuration.amq_adapter.service_name)
|
||||
self.amq_service_message_handler = AMQServiceMessageHandler(self.http_adapter, self.tracer, self.message_factory)
|
||||
self.amq_service_message_handler = None
|
||||
self.service_message_factory = CleverMicroServiceMessageFactory(
|
||||
self.amq_configuration.amq_adapter.service_name
|
||||
)
|
||||
self.register_routes()
|
||||
# self.init(loop) will initialize RabbitMQ connection
|
||||
loop.run_until_complete(self.init(loop, service_adapter))
|
||||
# loop.run_until_complete(consume(loop))
|
||||
# self.consumer_thread = asyncio.create_task(self.rabbit_mq_consumer.channel.start_consuming())
|
||||
# self.rabbit_mq_consumer.request_routes()
|
||||
|
||||
logging.info("***********************************************")
|
||||
logging.info("****** AMQ Service Init Complete *********")
|
||||
logging.info("***********************************************")
|
||||
loop.run_forever()
|
||||
|
||||
async def init(self, loop, service_adapter):
|
||||
exchange: AbstractRobustExchange = await self.rabbit_mq_consumer.init(loop)
|
||||
self.amq_service_message_handler = AMQServiceMessageHandler(
|
||||
service_adapter,
|
||||
self.tracer,
|
||||
self.message_factory,
|
||||
exchange,
|
||||
self.rabbit_mq_consumer.channel,
|
||||
self.amq_configuration.dispatch.rabbit_mq_url,
|
||||
output_dir=self.amq_configuration.dispatch.download_dir
|
||||
)
|
||||
await self.register_routes()
|
||||
await self.rabbit_mq_consumer.request_routes()
|
||||
|
||||
def cleanup(self):
|
||||
logging.info(" [DBG] AMQService.cleanup() rmqc=%s", self.rabbit_mq_consumer)
|
||||
if self.rabbit_mq_consumer:
|
||||
self.rabbit_mq_consumer.cleanup()
|
||||
|
||||
def reinitialize_if_disconnected(self):
|
||||
if not self.rabbit_mq_consumer.channel or not self.rabbit_mq_consumer.channel.is_open:
|
||||
async def reinitialize_if_disconnected(self):
|
||||
if not self.rabbit_mq_consumer.channel or self.rabbit_mq_consumer.channel.is_closed:
|
||||
self.rabbit_mq_consumer = RabbitMQConsumer(self.amq_configuration)
|
||||
self.register_routes()
|
||||
self.rabbit_mq_consumer.request_routes()
|
||||
await self.rabbit_mq_consumer.init(loop=None)
|
||||
await self.register_routes()
|
||||
await self.rabbit_mq_consumer.request_routes()
|
||||
|
||||
def register_routes(self):
|
||||
async def register_routes(self) -> AbstractRobustExchange | None:
|
||||
route_mapping = self.amq_configuration.amq_adapter.route_mapping
|
||||
if AMQRouteFactory.validate(route_mapping):
|
||||
try:
|
||||
self.rabbit_mq_consumer.register_inbound_routes(
|
||||
self.amq_service_message_handler.inbound_message_callback_http
|
||||
await self.rabbit_mq_consumer.register_inbound_routes(
|
||||
route_mapping, self.amq_service_message_handler.inbound_message_callback
|
||||
)
|
||||
self.rabbit_mq_consumer.register_rpc_handler(
|
||||
return await self.rabbit_mq_consumer.register_rpc_handler(
|
||||
self.amq_service_message_handler.reply_received_callback
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(e)
|
||||
elif route_mapping:
|
||||
logging.info(" [W] route [%s] failed validation, no routes were registered.", route_mapping)
|
||||
return None
|
||||
|
||||
def send_message(self, message: AMQMessage, future: Future):
|
||||
self.amq_service_message_handler.outstanding[str(message.id)] = future
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
import requests
|
||||
|
||||
|
||||
def delete_queues_with_prefix_http(host='localhost', port=15672, username='springuser', password='TheCleverWho', prefix='cm-'):
|
||||
"""
|
||||
Delete queues using RabbitMQ HTTP API.
|
||||
"""
|
||||
base_url = f"http://{host}:{port}/api/queues/%2F"
|
||||
|
||||
try:
|
||||
response = requests.get(base_url, auth=(username, password))
|
||||
response.raise_for_status()
|
||||
|
||||
queues = [q['name'] for q in response.json() if q['name'].startswith(prefix)]
|
||||
|
||||
if not queues:
|
||||
print(f"No queues found with prefix '{prefix}'")
|
||||
return
|
||||
|
||||
print(f"Found {len(queues)} queues with prefix '{prefix}':")
|
||||
for queue in queues:
|
||||
print(f" - {queue}")
|
||||
|
||||
for queue in queues:
|
||||
try:
|
||||
delete_url = f"{base_url}/{queue}"
|
||||
response = requests.delete(delete_url, auth=(username, password))
|
||||
response.raise_for_status()
|
||||
print(f"Successfully deleted queue: {queue}")
|
||||
except Exception as e:
|
||||
print(f"Failed to delete queue {queue}: {str(e)}")
|
||||
|
||||
print("Queue deletion process completed.")
|
||||
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {str(e)}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
delete_queues_with_prefix_http()
|
||||
@@ -0,0 +1,50 @@
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
|
||||
import python_multipart
|
||||
from fastapi import UploadFile
|
||||
from python_multipart import multipart
|
||||
from python_multipart.multipart import Field, File
|
||||
|
||||
from amqp.model.model import AMQMessage
|
||||
|
||||
|
||||
class CleverMultiPartParser:
|
||||
def __init__(self, message: AMQMessage):
|
||||
self.form_data = {}
|
||||
self.is_multipart = False
|
||||
self.is_json = False
|
||||
self.is_valid = True
|
||||
# Convert each list of strings to a single string joined by newline, then to bytes
|
||||
headers = {key: "\n".join(value).encode("utf-8") for key, value in message.headers.items()}
|
||||
content_type: str | bytes | None = headers.get("Content-Type")
|
||||
content_type, params = multipart.parse_options_header(content_type)
|
||||
if content_type.decode('utf-8') in ["application/octet-stream",
|
||||
"multipart/form-data",
|
||||
"application/x-www-form-urlencoded",
|
||||
"application/x-url-encoded"]:
|
||||
try:
|
||||
python_multipart.parse_form(headers, io.BytesIO(message.body()), self.on_field, self.on_file)
|
||||
self.is_multipart = True
|
||||
except Exception as e:
|
||||
logging.error(f"Error parsing multipart/form-data: {e}. Trying parsing as JSON.")
|
||||
try:
|
||||
self.form_data = json.loads(message.body())
|
||||
self.is_json = True
|
||||
except Exception as jsonerror:
|
||||
logging.error(f"Error parsing JSON: {jsonerror}")
|
||||
self.is_valid = False
|
||||
else:
|
||||
self.form_data = message.body()
|
||||
|
||||
def on_field(self, field: Field):
|
||||
print(field)
|
||||
self.form_data[field.field_name.decode('utf-8')] = field.value
|
||||
print(self.form_data)
|
||||
|
||||
def on_file(self, file: File):
|
||||
print(file)
|
||||
self.form_data[file.field_name.decode('utf-8')] = \
|
||||
UploadFile(file.file_object, size=file.size, filename=file.file_name)
|
||||
print(self.form_data)
|
||||
@@ -1,23 +1,31 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
import concurrent.futures
|
||||
import os
|
||||
import traceback
|
||||
from asyncio import Future
|
||||
from typing import Dict
|
||||
|
||||
import pika
|
||||
from aio_pika import IncomingMessage, Message
|
||||
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustExchange, AbstractRobustChannel
|
||||
|
||||
from opentelemetry import trace
|
||||
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
|
||||
from opentelemetry.trace import Tracer, TraceState
|
||||
from pika.channel import Channel
|
||||
|
||||
from amqp.adapter.amq_message_factory import AMQMessageFactory
|
||||
from amqp.adapter.amq_response_factory import AMQResponseFactory
|
||||
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
|
||||
from amqp.model.model import AMQResponse, AMQMessage
|
||||
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
|
||||
from amqp.adapter.file_downloader import on_message
|
||||
from amqp.model.model import AMQResponse, AMQMessage, MultipartAMQMessage
|
||||
from amqp.router.serializer import map_as_string
|
||||
|
||||
parallel_workers = os.getenv("PARALLEL_WORKERS", default=10)
|
||||
|
||||
# Shared thread pool for parallel execution
|
||||
executor = concurrent.futures.ThreadPoolExecutor(max_workers=parallel_workers)
|
||||
|
||||
def collect_trace_info():
|
||||
trace_map = {}
|
||||
@@ -37,66 +45,95 @@ def set_text(carrier, key, value):
|
||||
|
||||
|
||||
class AMQServiceMessageHandler:
|
||||
def __init__(self, service_adapter: CleverThisServiceAdapter, tracer: Tracer, message_factory: AMQMessageFactory):
|
||||
self.http_adapter: CleverThisServiceAdapter = service_adapter
|
||||
def __init__(self,
|
||||
service_adapter: CleverThisServiceAdapter,
|
||||
tracer: Tracer,
|
||||
message_factory: AMQMessageFactory,
|
||||
reply_to_exchange: AbstractRobustExchange,
|
||||
channel: AbstractRobustChannel,
|
||||
rabbit_mq_url: str,
|
||||
output_dir: str):
|
||||
self.service_adapter: CleverThisServiceAdapter = service_adapter
|
||||
self.tracer = tracer
|
||||
self.message_factory: AMQMessageFactory = message_factory
|
||||
self.reply_to_exchange: AbstractRobustExchange = reply_to_exchange
|
||||
self.channel: AbstractRobustChannel = channel
|
||||
self.rabbit_mq_url = rabbit_mq_url
|
||||
self.output_dir = output_dir
|
||||
# Dictionary to store outstanding Futures
|
||||
self.outstanding: Dict[str, asyncio.Future] = {}
|
||||
|
||||
def inbound_message_callback_http(self,
|
||||
channel: Channel,
|
||||
method: pika.spec.Basic.Deliver,
|
||||
properties: pika.spec.BasicProperties,
|
||||
body: bytes):
|
||||
async def inbound_message_callback(self, message: AbstractIncomingMessage):
|
||||
"""
|
||||
Receives the message as bytes from RabbitMQ queue and deserializes it to the AMQMessage
|
||||
Ensures Jaeger trace-id propagation.
|
||||
In loose coupling mode, invokes appropriate adapter (in v0.1 only HTTPAdapter is supported).
|
||||
In tight coupling mode, invokes custom handler/mapper/adapter method that
|
||||
:param body:
|
||||
:param properties:
|
||||
:param method:
|
||||
:param channel:
|
||||
Invokes custom handler/mapper/adapter method that corresponds to this message
|
||||
:param message:
|
||||
:return:
|
||||
"""
|
||||
start = time.time()
|
||||
delivery_tag = method.delivery_tag
|
||||
debug = f" [DBG] ######[{delivery_tag}] (AMQ MESSAGE), env: {method}, props: {properties}"
|
||||
logging.info(debug)
|
||||
delivery_tag = message.delivery_tag
|
||||
debug = f" [DBG] ######[{delivery_tag}] (AMQ MESSAGE), env: {message.headers}, props: {message.properties}"
|
||||
logging.debug(debug)
|
||||
addressing: int = message.headers.get('clevermicro.addressing', 0)
|
||||
|
||||
amq_message: AMQMessage = AMQMessageFactory.from_bytes(body)
|
||||
logging.info(" [*] ######[%s] AMQ Message received, ctag=%s, msgId=%s, traceInfo=%s",
|
||||
delivery_tag, method.consumer_tag, amq_message.id, map_as_string(amq_message.trace_info))
|
||||
|
||||
propagator = TraceContextTextMapPropagator()
|
||||
context = propagator.extract(carrier=amq_message.trace_info, context=trace.context_api.get_current())
|
||||
self.tracer.start_span(name="AMQ-Adapter-Consumer", context=context)
|
||||
|
||||
logging.info(" [*] CTX=%s", context)
|
||||
# current_span = self.tracer.start_span("AMQ-Adapter.Processing.Request", parent=parent_span)
|
||||
# context_with_span = trace.set_span_in_context(current_span, context)
|
||||
# logging.info(" [*] CTX2=%s", context_with_span)
|
||||
# propagator.inject(context_with_span, amq_message.trace_info, self.set_text)
|
||||
# logging.info(" [***] NEW TI=%s CTX=%s SCOPE=%s",
|
||||
# map_as_string(amq_message.trace_info), context_with_span, context)
|
||||
|
||||
if amq_message.magic == AMQMessageFactory.MAGIC_BYTE_HTTP_REQUEST:
|
||||
response: AMQResponse = asyncio.run(self.http_adapter.on_message(amq_message))
|
||||
self.send_reply(properties.reply_to, delivery_tag, channel, response)
|
||||
amq_message: AMQMessage | None = None
|
||||
if addressing == 0:
|
||||
amq_message = AMQMessageFactory.from_bytes(message.body)
|
||||
else:
|
||||
logging.error(" [E] Unknown MAGIC value: [%s], don't know how to interpret the message",
|
||||
amq_message.magic)
|
||||
loop = asyncio.get_event_loop()
|
||||
amq_message = await AMQMessageFactory.from_stream(message.body, rabbitmq_url=self.rabbit_mq_url, loop=loop)
|
||||
logging.info(f" [DBG] ######[{delivery_tag}] AMQ Message from stream: {amq_message}")
|
||||
if amq_message is not None:
|
||||
extra_data = await on_message(
|
||||
message=message,
|
||||
json_data=amq_message.body(),
|
||||
rabbitmq_url=self.rabbit_mq_url,
|
||||
loop=loop,
|
||||
output_dir=self.output_dir)
|
||||
logging.info(f" [x] ######[{delivery_tag}] Multipart download result: {extra_data}")
|
||||
amq_message = MultipartAMQMessage(amq_message, extra_data)
|
||||
|
||||
if amq_message is not None:
|
||||
logging.info(" [*] ######[%s] AMQ Message received, ctag=%s, msgId=%s, traceInfo=%s",
|
||||
delivery_tag, message.consumer_tag, amq_message.id, map_as_string(amq_message.trace_info))
|
||||
|
||||
propagator = TraceContextTextMapPropagator()
|
||||
context = propagator.extract(carrier=amq_message.trace_info, context=trace.context_api.get_current())
|
||||
self.tracer.start_span(name="AMQ-Adapter-Consumer", context=context)
|
||||
|
||||
logging.info(" [*] CTX=%s", context)
|
||||
# current_span = self.tracer.start_span("AMQ-Adapter.Processing.Request", parent=parent_span)
|
||||
# context_with_span = trace.set_span_in_context(current_span, context)
|
||||
# logging.info(" [*] CTX2=%s", context_with_span)
|
||||
# propagator.inject(context_with_span, amq_message.trace_info, self.set_text)
|
||||
# logging.info(" [***] NEW TI=%s CTX=%s SCOPE=%s",
|
||||
# map_as_string(amq_message.trace_info), context_with_span, context)
|
||||
|
||||
if amq_message.magic == AMQMessageFactory.MAGIC_BYTE_HTTP_REQUEST:
|
||||
future = executor.submit(self.service_adapter.sync_on_message, amq_message)
|
||||
# response: AMQResponse = asyncio.run(self.http_adapter.on_message(amq_message))
|
||||
# Wait and send it back to the reply queue
|
||||
try:
|
||||
response: AMQResponse = future.result() # Block until the result is available
|
||||
await self.send_reply(message.reply_to, delivery_tag, response)
|
||||
except Exception as e:
|
||||
logging.info(f"[E] Error processing message: {e}")
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f"[E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
|
||||
await message.nack(requeue=False)
|
||||
return
|
||||
else:
|
||||
logging.error(" [E] Unknown MAGIC value: [%s], don't know how to interpret the message",
|
||||
amq_message.magic)
|
||||
await message.nack(requeue=False)
|
||||
return
|
||||
else:
|
||||
logging.error(" [E] ######[%s] No valid AMQ Message present.", delivery_tag)
|
||||
|
||||
duration = time.time() - start
|
||||
logging.info(" [x] ######[%s] Done, single ACK: %sms.", delivery_tag, duration)
|
||||
channel.basic_ack(delivery_tag, False)
|
||||
# if "___NACK___" in str(amq_message.body()):
|
||||
# logging.info(" [x] ######[%s] Done, single NACK: %sms.", delivery_tag, duration)
|
||||
# channel.basic_nack(delivery_tag, False, False)
|
||||
# else:
|
||||
# logging.info(" [x] ######[%s] Done, single ACK: %sms.", delivery_tag, duration)
|
||||
# channel.basic_ack(delivery_tag, False)
|
||||
await message.ack(multiple=False)
|
||||
|
||||
def on_http_response(self, delivery, deliveryTag, channel, http_response):
|
||||
"""
|
||||
@@ -119,53 +156,47 @@ class AMQServiceMessageHandler:
|
||||
collect_trace_info()
|
||||
)
|
||||
try:
|
||||
self.send_reply(delivery.properties.reply_to, deliveryTag, channel, response)
|
||||
self.send_reply(delivery.properties.reply_to, deliveryTag, response)
|
||||
except Exception as e:
|
||||
tb = traceback.extract_tb(sys.exc_info()[2])[-1]
|
||||
logging.error(f"[E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
|
||||
raise RuntimeError(e)
|
||||
|
||||
def send_reply(self, reply_to: str, delivery_tag: str, channel: Channel, response: AMQResponse):
|
||||
async def send_reply(self, reply_to: str, delivery_tag: int | None, response: AMQResponse):
|
||||
"""
|
||||
The Service side code. When the CleverXXX service has output ready, in AMQResponse object,
|
||||
call this method to return the reply back to the caller
|
||||
:param reply_to:
|
||||
:param delivery_tag:
|
||||
:param channel:
|
||||
:param response:
|
||||
:return:
|
||||
"""
|
||||
if reply_to:
|
||||
logging.info(" [*] ######[%s] AMQ Message Reply, To=%s, msgId=%s",
|
||||
delivery_tag, reply_to, response.id)
|
||||
reply_props = pika.BasicProperties(
|
||||
delivery_tag, reply_to, response.id)
|
||||
pika_message: Message = Message(
|
||||
body=AMQResponseFactory.serialize(response),
|
||||
correlation_id=response.id,
|
||||
content_type=response.content_type[0]
|
||||
if isinstance(response.content_type, list) else response.content_type
|
||||
)
|
||||
channel.basic_publish(
|
||||
exchange=AMQ_REPLY_TO_EXCHANGE,
|
||||
routing_key=reply_to,
|
||||
body=AMQResponseFactory.serialize(response),
|
||||
properties=reply_props,
|
||||
await self.reply_to_exchange.publish(
|
||||
message=pika_message,
|
||||
routing_key=reply_to
|
||||
)
|
||||
|
||||
def reply_received_callback(self,
|
||||
channel: Channel,
|
||||
method: pika.spec.Basic.Deliver,
|
||||
properties: pika.spec.BasicProperties,
|
||||
body: bytes):
|
||||
async def reply_received_callback(self, message: IncomingMessage):
|
||||
"""
|
||||
The producer/client (API Gateway) facing response_message recipient - must match the received response with
|
||||
the original request and produce the HTTP reply.
|
||||
Note: this code runs on the API Gateway side of RabbitMQ
|
||||
:param body:
|
||||
:param properties:
|
||||
:param method:
|
||||
:param channel:
|
||||
:param message:
|
||||
:return:
|
||||
"""
|
||||
reply_id = properties.correlation_id
|
||||
content_type = properties.content_type
|
||||
delivery_tag = method.delivery_tag
|
||||
debug = f" [DBG] ######[{delivery_tag}] (RPC REPLY) env: {method}, props: {properties}, BODY: {body.decode()}"
|
||||
reply_id = message.correlation_id
|
||||
content_type = message.content_type
|
||||
delivery_tag = message.delivery_tag
|
||||
debug = f" [DBG] ######[{delivery_tag}] (RPC REPLY) env: {message.headers}, props: {message.properties}, BODY: {message.body.decode()}"
|
||||
logging.info(debug)
|
||||
# TODO - figure out implementation of this part, as this is sending reply back either to a service,
|
||||
# or to the user
|
||||
@@ -183,8 +214,8 @@ class AMQServiceMessageHandler:
|
||||
# internal, service-2-service communication, and the proper response type will be decided than
|
||||
# reply.response().set_result(Response.ok(result.body, content_type).build())
|
||||
if future:
|
||||
response: AMQResponse = AMQResponseFactory.from_bytes(body)
|
||||
response: AMQResponse = AMQResponseFactory.from_bytes(message.body)
|
||||
# Complete the Future with the response data
|
||||
future.set_result(response.body())
|
||||
|
||||
channel.basic_ack(delivery_tag, False)
|
||||
await message.ack(multiple=False)
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
import json
|
||||
import struct
|
||||
import io
|
||||
|
||||
|
||||
class StreamParser:
|
||||
"""
|
||||
Parses input stream that contains metadata followed by binary data.
|
||||
First 4 bytes are metadata length, followed by metadata in JSON format, followed by binary data.
|
||||
The StreamParser accumulates the binary data from individual chunks and writes it to a file.
|
||||
This can be used to parse a stream that contains a large file upload.
|
||||
"""
|
||||
def __init__(self):
|
||||
self.buffer = io.BytesIO() # Buffer to accumulate chunks
|
||||
self.metadata_length = None # Length of the metadata
|
||||
self.metadata = None # Parsed JSON metadata
|
||||
self.file_writer = None # File writer for the remaining data
|
||||
seek_to = 0
|
||||
|
||||
def process_chunk(self, chunk: bytes):
|
||||
"""
|
||||
Process a chunk of data and handle it according to the rules.
|
||||
"""
|
||||
self.buffer.write(chunk)
|
||||
available_data = self.buffer.tell()
|
||||
# If metadata length is not yet known, try to read it
|
||||
if self.metadata_length is None and available_data >= 4:
|
||||
self.buffer.seek(0)
|
||||
self.metadata_length = struct.unpack('>I', self.buffer.read(4))[0]
|
||||
self.seek_to = self.buffer.tell()
|
||||
|
||||
# If metadata length is known but metadata is not yet parsed, try to parse it
|
||||
if self.metadata_length is not None and self.metadata is None:
|
||||
if available_data >= 4 + self.metadata_length:
|
||||
self.buffer.seek(self.seek_to)
|
||||
metadata_bytes = self.buffer.read(self.metadata_length)
|
||||
self.metadata = json.loads(metadata_bytes.decode('utf-8'))
|
||||
self.seek_to = self.buffer.tell()
|
||||
print("Metadata:", self.metadata)
|
||||
|
||||
|
||||
# If metadata is parsed, write the remaining data to the file
|
||||
if self.metadata is not None:
|
||||
self.buffer.seek(self.seek_to)
|
||||
remaining_data = self.buffer.read()
|
||||
self.seek_to = self.buffer.tell()
|
||||
if remaining_data:
|
||||
if self.file_writer is None:
|
||||
# Prepare to write the remaining data to a file
|
||||
self.file_writer = open(self.metadata.get("filename", "output_file"), "wb")
|
||||
self.file_writer.write(remaining_data)
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Close the file writer and clean up.
|
||||
"""
|
||||
if self.file_writer:
|
||||
self.file_writer.close()
|
||||
self.buffer.close()
|
||||
@@ -0,0 +1,141 @@
|
||||
import unittest
|
||||
import json
|
||||
import struct
|
||||
import os
|
||||
from unittest.mock import patch, mock_open
|
||||
|
||||
from amqp.service.stream_parser import StreamParser
|
||||
|
||||
|
||||
class TestStreamParser(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""
|
||||
Set up the StreamParser instance for testing.
|
||||
"""
|
||||
self.parser = StreamParser()
|
||||
|
||||
def tearDown(self):
|
||||
"""
|
||||
Clean up after each test.
|
||||
"""
|
||||
if hasattr(self.parser, 'file_writer') and self.parser.file_writer:
|
||||
self.parser.file_writer.close()
|
||||
if os.path.exists("test_output_file"):
|
||||
os.remove("test_output_file")
|
||||
|
||||
def test_process_chunk_with_metadata_and_data(self):
|
||||
"""
|
||||
Test processing a chunk that contains metadata and data.
|
||||
"""
|
||||
# Create test data
|
||||
metadata = {"filename": "test_output_file", "description": "Test file"}
|
||||
metadata_json = json.dumps(metadata).encode('utf-8')
|
||||
metadata_length = struct.pack('>I', len(metadata_json))
|
||||
data = b"some binary data"
|
||||
|
||||
# Combine metadata length, metadata, and data into a single chunk
|
||||
chunk = metadata_length + metadata_json + data
|
||||
|
||||
# Process the chunk
|
||||
self.parser.process_chunk(chunk)
|
||||
self.parser.close()
|
||||
|
||||
# Verify metadata
|
||||
self.assertEqual(self.parser.metadata, metadata)
|
||||
self.assertEqual(self.parser.metadata_length, len(metadata_json))
|
||||
|
||||
# Verify file content
|
||||
with open("test_output_file", "rb") as f:
|
||||
file_content = f.read()
|
||||
self.assertEqual(file_content, data)
|
||||
|
||||
def test_process_chunk_in_parts(self):
|
||||
"""
|
||||
Test processing chunks in multiple parts.
|
||||
"""
|
||||
# Create test data
|
||||
metadata = {"filename": "test_output_file", "description": "Test file"}
|
||||
metadata_json = json.dumps(metadata).encode('utf-8')
|
||||
metadata_length = struct.pack('>I', len(metadata_json))
|
||||
data = b"some binary data"
|
||||
|
||||
# Split the chunk into parts
|
||||
chunk1 = metadata_length # First 4 bytes (metadata length)
|
||||
chunk2 = metadata_json # Metadata JSON
|
||||
chunk3 = data # Remaining data
|
||||
|
||||
# Process chunks one by one
|
||||
self.parser.process_chunk(chunk1)
|
||||
self.parser.process_chunk(chunk2)
|
||||
self.parser.process_chunk(chunk3)
|
||||
self.parser.close()
|
||||
|
||||
# Verify metadata
|
||||
self.assertEqual(self.parser.metadata, metadata)
|
||||
self.assertEqual(self.parser.metadata_length, len(metadata_json))
|
||||
|
||||
# Verify file content
|
||||
with open("test_output_file", "rb") as f:
|
||||
file_content = f.read()
|
||||
self.assertEqual(file_content, data)
|
||||
|
||||
def test_process_chunk_without_metadata(self):
|
||||
"""
|
||||
Test processing a chunk without metadata.
|
||||
"""
|
||||
# Create test data (only metadata length and metadata, no additional data)
|
||||
metadata = {"filename": "distinct_test_output_file", "description": "Test file"}
|
||||
metadata_json = json.dumps(metadata).encode('utf-8')
|
||||
metadata_length = struct.pack('>I', len(metadata_json))
|
||||
|
||||
# Combine metadata length and metadata into a single chunk
|
||||
chunk = metadata_length + metadata_json
|
||||
|
||||
# Process the chunk
|
||||
self.parser.process_chunk(chunk)
|
||||
self.parser.close()
|
||||
|
||||
# Verify metadata
|
||||
self.assertEqual(self.parser.metadata, metadata)
|
||||
self.assertEqual(self.parser.metadata_length, len(metadata_json))
|
||||
|
||||
# Verify no data was written to the file
|
||||
self.assertFalse(os.path.exists("distinct_test_output_file"))
|
||||
|
||||
def test_process_chunk_with_invalid_metadata(self):
|
||||
"""
|
||||
Test processing a chunk with invalid metadata.
|
||||
"""
|
||||
# Create invalid metadata (not JSON)
|
||||
metadata_length = struct.pack('>I', 10)
|
||||
invalid_metadata = b"invalid_json"
|
||||
|
||||
# Combine metadata length and invalid metadata into a single chunk
|
||||
chunk = metadata_length + invalid_metadata
|
||||
|
||||
# Process the chunk and expect an exception
|
||||
with self.assertRaises(json.JSONDecodeError):
|
||||
self.parser.process_chunk(chunk)
|
||||
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open)
|
||||
def test_file_writing(self, mock_file):
|
||||
"""
|
||||
Test file writing functionality using a mock file.
|
||||
"""
|
||||
# Create test data
|
||||
metadata = {"filename": "test_output_file", "description": "Test file"}
|
||||
metadata_json = json.dumps(metadata).encode('utf-8')
|
||||
metadata_length = struct.pack('>I', len(metadata_json))
|
||||
data = b"some binary data"
|
||||
|
||||
# Combine metadata length, metadata, and data into a single chunk
|
||||
chunk = metadata_length + metadata_json + data
|
||||
|
||||
# Process the chunk
|
||||
self.parser.process_chunk(chunk)
|
||||
|
||||
# Verify the file was opened and written to
|
||||
mock_file.assert_called_once_with("test_output_file", "wb")
|
||||
mock_file().write.assert_called_once_with(data)
|
||||
@@ -1,5 +1,4 @@
|
||||
from opentelemetry import metrics, trace
|
||||
#from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
|
||||
from opentelemetry.sdk.metrics import MeterProvider
|
||||
from opentelemetry.sdk.trace import TracerProvider, Tracer
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
||||
|
||||
@@ -0,0 +1,287 @@
|
||||
import logging
|
||||
import pathlib
|
||||
import re
|
||||
from typing import Tuple, AnyStr, Any
|
||||
from aiohttp import ClientSession, ClientResponse
|
||||
from starlette.datastructures import UploadFile, Headers
|
||||
import jwt
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import core.deps
|
||||
from amqp.adapter.amq_response_factory import AMQResponseFactory
|
||||
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, call_cleverthis_get_api, \
|
||||
call_cleverthis_post_put_api, parse_request_body
|
||||
from amqp.service.multipart_parser import CleverMultiPartParser
|
||||
from amqp.config.amq_configuration import AMQConfiguration
|
||||
from amqp.model.model import AMQMessage, AMQResponse
|
||||
|
||||
import rest.v0.endpoints.benchmark as benchmark
|
||||
import rest.v0.endpoints.jobs as jobs
|
||||
import rest.v0.endpoints.unstructured as unstructured
|
||||
from amqp.service.amq_service import AMQService
|
||||
from core.configs import settings
|
||||
from schemas.user_schema import UserSchema
|
||||
|
||||
|
||||
BENCHMARK = '/benchmark'
|
||||
JOBS = '/jobs/'
|
||||
JOB_DETAILS = '/jobs/details/'
|
||||
UNSTRUCTURED = '/unstructured'
|
||||
|
||||
GET_BENCHMARK_KG = r".*/benchmark/(?P<job_id>[^\?]+)\?index=(?P<index>\d+)"
|
||||
GET_JOB_STATUS = r".*/jobs/(?P<job_id>.*)"
|
||||
GET_KG = r".*/unstructured/(?P<job_id>.*)"
|
||||
PUT_BENCHMARK_KG = r".*/benchmark/(?P<job_id>.*)"
|
||||
DUMMY = r"(?P<job_id>.*)"
|
||||
|
||||
|
||||
#
|
||||
# Following methods are wrappers to actual CleverSwarm endpoints, so I can refer to the CleverSwarm method and
|
||||
# pass it as input parameter using general Callable[[Any, UserSchema], str] type
|
||||
#
|
||||
async def cleverswarm_create_benchmark_job(args: Any, authenticated_user: UserSchema) -> str:
|
||||
return str(await benchmark.create_benchmark_job(**args, authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
async def cleverswarm_create_extract_with_ontology(args: Any, authenticated_user: UserSchema) -> str:
|
||||
return str(await unstructured.create_extract_with_ontology(**args, authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
async def cleverswarm_create_extract_with_wildcards(args: Any, authenticated_user: UserSchema) -> str:
|
||||
return str(await unstructured.create_extract_with_wildcards(**args, authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
async def cleverswarm_update_benchmark_job(args: Any, authenticated_user: UserSchema) -> str:
|
||||
return str(await benchmark.update_benchmark_job(**args, authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
async def cleverswarm_retry_job(args: Any, authenticated_user: UserSchema) -> str:
|
||||
return str(await jobs.retry_job(args[0], authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
async def cleverswarm_delete_job(args: Any, authenticated_user: UserSchema) -> str:
|
||||
return str(await jobs.delete_job(args[0], authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
async def cleverswarm_get_benchmark_kg(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
|
||||
return str(await benchmark.get_benchmark_kg(args[0], args[1], authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
async def cleverswarm_get_job_status(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
|
||||
return str(await jobs.get_job_status(args[0], authenticated_user=authenticated_user))
|
||||
|
||||
async def cleverswarm_get_job_details(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
|
||||
return str(await jobs.get_job_details(args[0], authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
async def cleverswarm_jobs_get_requests(args: Tuple[AnyStr | Any, ...],
|
||||
authenticated_user: UserSchema) -> str:
|
||||
return str(await jobs.get_requests(authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
async def cleverswarm_get_kg(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
|
||||
return str(await unstructured.get_kg(args[0], authenticated_user=authenticated_user))
|
||||
|
||||
|
||||
class CleverSwarmAmqpAdapter(CleverThisServiceAdapter):
|
||||
def __init__(self, amq_configuration: AMQConfiguration, session: ClientSession = None):
|
||||
super().__init__(amq_configuration, session)
|
||||
amq_service: AMQService = AMQService(amq_configuration,self)
|
||||
amq_service.rabbit_mq_consumer.channel.start_consuming()
|
||||
|
||||
async def get_current_user(self, token: str) -> UserSchema:
|
||||
if token == 'DEBUG_NO_AUTH':
|
||||
# Payload with username and expiration
|
||||
payload = {
|
||||
"sub": "0", # Standard claim for subject (user)
|
||||
"username": "user", # Custom claim (optional)
|
||||
"iat": datetime.utcnow(), # Issued at time
|
||||
"exp": datetime.utcnow() + timedelta(days=365), # Expires in 1 year
|
||||
}
|
||||
# Generate the JWT token
|
||||
token = jwt.encode(payload, settings.JWT_SECRET, algorithm=settings.ALGORITHM)
|
||||
|
||||
return await core.deps.get_current_user(token)
|
||||
# ==================================================================================================
|
||||
# ===================== C l e v e r S w a r m A P I m e t h o d s ============================
|
||||
# ==================================================================================================
|
||||
|
||||
async def get(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
|
||||
"""
|
||||
Handles the GET requests for the CleverSwarm API.
|
||||
:param message: message from the AMQP that contains the GET request
|
||||
:param auth_user: user authenticated by the API Gateway, as UserSchem
|
||||
:return: AMQResponse to be sent back to the API Gateway via AMQP
|
||||
"""
|
||||
if BENCHMARK in message.path:
|
||||
return await call_cleverthis_get_api(
|
||||
message,
|
||||
GET_BENCHMARK_KG,
|
||||
cleverswarm_get_benchmark_kg,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
|
||||
if JOBS in message.path:
|
||||
if JOB_DETAILS in message.path:
|
||||
return await call_cleverthis_get_api(
|
||||
message,
|
||||
DUMMY,
|
||||
cleverswarm_get_job_status,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
return await call_cleverthis_get_api(
|
||||
message,
|
||||
GET_JOB_STATUS,
|
||||
cleverswarm_get_job_status,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
|
||||
if message.path.endswith('/jobs'):
|
||||
return await call_cleverthis_get_api(
|
||||
message,
|
||||
DUMMY,
|
||||
cleverswarm_jobs_get_requests,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
|
||||
if UNSTRUCTURED in message.path:
|
||||
return await call_cleverthis_get_api(
|
||||
message,
|
||||
GET_KG,
|
||||
cleverswarm_get_kg,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
|
||||
# unmatched path - try to invoke the service directly on this path
|
||||
url = f"http://{self.service_host}:{self.service_port}{message.path}"
|
||||
headers = {**message.headers, **message.trace_info}
|
||||
logging.info(f" [*] SEND HEADERS out: {', '.join(f'{k}={v}' for k, v in headers.items())}, REQ={url}")
|
||||
if self.session:
|
||||
_http_resp: ClientResponse = await self.session.get(url, headers=headers)
|
||||
return AMQResponseFactory.of(message.id, _http_resp.status, _http_resp.content_type,
|
||||
await _http_resp.read(),
|
||||
message.trace_info)
|
||||
return AMQResponseFactory.of(message.id, 404, "text/plain",
|
||||
str(f"Destination location [{url}] does not exists").encode('utf-8'),
|
||||
message.trace_info)
|
||||
|
||||
async def put(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
|
||||
"""
|
||||
Handles the PUT requests for the CleverSwarm API.
|
||||
:param message: message from the AMQP that contains the PUT request
|
||||
:param auth_user: user authenticated by the API Gateway, as UserSchem
|
||||
:return: AMQResponse to be sent back to the API Gateway via AMQP
|
||||
"""
|
||||
parser: CleverMultiPartParser = CleverMultiPartParser(message)
|
||||
if BENCHMARK in message.path:
|
||||
match = re.search(PUT_BENCHMARK_KG, message.path)
|
||||
if match:
|
||||
parser.form_data['job_id'] = match.group('job_id')
|
||||
return await call_cleverthis_post_put_api(
|
||||
message,
|
||||
parser.form_data,
|
||||
cleverswarm_update_benchmark_job,
|
||||
201,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
|
||||
if JOBS in message.path:
|
||||
return await call_cleverthis_get_api(
|
||||
message,
|
||||
GET_JOB_STATUS,
|
||||
cleverswarm_retry_job,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
url = f"http://{self.service_host}:{self.service_port}{message.path}"
|
||||
return await self.handle_possible_form(
|
||||
self.session.put(
|
||||
url,
|
||||
data=message.body,
|
||||
headers={**message.headers, **message.trace_info},
|
||||
),
|
||||
auth_user
|
||||
) if self.session else AMQResponseFactory.of(message.id, 404, "text/plain",
|
||||
str(f"Destination location [{url}] does not exists").encode('utf-8'),
|
||||
message.trace_info)
|
||||
|
||||
async def post(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
|
||||
"""
|
||||
Handles the POST requests for the CleverSwarm API.
|
||||
:param message: message from the AMQP that contains the POST request, including body and headers
|
||||
:param auth_user: user authenticated by the API Gateway, as UserSchem
|
||||
:return: AMQResponse to be sent back to the API Gateway via AMQP
|
||||
"""
|
||||
_args = parse_request_body(message)
|
||||
args = {}
|
||||
if isinstance(_args, dict):
|
||||
for key, file in _args['files'].items():
|
||||
if len(file) > 0:
|
||||
filepath = message.extra_data[key]
|
||||
ufile = UploadFile(
|
||||
file=pathlib.Path(filepath).open("rb"),
|
||||
size=file[0]['size'],
|
||||
filename=file[0]['filename'],
|
||||
headers=Headers({"Content-Type": file[0]['mediaType']})
|
||||
)
|
||||
args[key] = ufile
|
||||
for key, obj in _args['form'].items():
|
||||
if key not in args:
|
||||
args[key] = obj[0] if isinstance(obj, list) else obj
|
||||
|
||||
if BENCHMARK in message.path:
|
||||
return await call_cleverthis_post_put_api(
|
||||
message,
|
||||
args,
|
||||
cleverswarm_create_benchmark_job,
|
||||
201,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
if UNSTRUCTURED in message.path:
|
||||
if 'with_ontology' in message.path:
|
||||
return await call_cleverthis_post_put_api(
|
||||
message,
|
||||
args,
|
||||
cleverswarm_create_extract_with_ontology,
|
||||
200,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
if 'with_wildcards' in message.path:
|
||||
return await call_cleverthis_post_put_api(
|
||||
message,
|
||||
args,
|
||||
cleverswarm_create_extract_with_wildcards,
|
||||
200,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
url = f"http://{self.service_host}:{self.service_port}{message.path}"
|
||||
return await self.handle_possible_form(
|
||||
self.session.post(
|
||||
url,
|
||||
data=message.body,
|
||||
headers={**message.headers, **message.trace_info},
|
||||
),
|
||||
auth_user
|
||||
) if self.session else AMQResponseFactory.of(message.id, 404, "text/plain",
|
||||
str(f"Destination location [{url}] does not exists").encode('utf-8'),
|
||||
message.trace_info)
|
||||
|
||||
async def handle_possible_form(self,
|
||||
message: AMQMessage,
|
||||
request_coroutine,
|
||||
auth_user: UserSchema) -> AMQResponse:
|
||||
return await request_coroutine()
|
||||
|
||||
async def head(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
|
||||
raise NotImplementedError
|
||||
|
||||
async def delete(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
|
||||
if JOBS in message.path:
|
||||
return await call_cleverthis_get_api(
|
||||
message,
|
||||
GET_JOB_STATUS,
|
||||
cleverswarm_retry_job,
|
||||
authenticated_user=auth_user
|
||||
)
|
||||
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,14 @@
|
||||
import getch
|
||||
from amqp.config.amq_configuration import AMQConfiguration
|
||||
|
||||
from clever_swarm_adapter import CleverSwarmAmqpAdapter
|
||||
|
||||
if __name__ == "__main__":
|
||||
adapter = CleverSwarmAmqpAdapter(AMQConfiguration('./application.properties.local'))
|
||||
|
||||
print("Press the space bar to exit...")
|
||||
while True:
|
||||
key = getch.getch() # Wait for a keypress
|
||||
if key == ' ': # Check if the key is the space bar
|
||||
print("Space bar pressed. Exiting...")
|
||||
break
|
||||
Reference in New Issue
Block a user