Add support for file streaming upload, update to async pika lib
This commit is contained in:
Stanislav Hejny
+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
|
||||
|
||||
Reference in New Issue
Block a user