feat/#1 - refactor code for better design and testability
This commit was merged in pull request #11.
This commit is contained in:
Stanislav Hejny
@@ -1,5 +1,3 @@
|
||||
import logging
|
||||
|
||||
from amqp.adapter.logging_utils import logging_error
|
||||
from amqp.model.model import AMQRoute
|
||||
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from asyncio import AbstractEventLoop
|
||||
|
||||
from aio_pika import Message
|
||||
from aio_pika.abc import AbstractRobustChannel, ExchangeType
|
||||
|
||||
from amqp.model.model import ScalingRequestAlert
|
||||
|
||||
|
||||
class ScaleRequestV1:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
serviceId: str,
|
||||
taskId: str,
|
||||
maxAvailability: int,
|
||||
currentAvailability: int,
|
||||
requestType: ScalingRequestAlert,
|
||||
):
|
||||
self.version = 1 # Version of the request, currently 1
|
||||
self.serviceId = serviceId
|
||||
self.taskId = taskId
|
||||
self.maxAvailability = maxAvailability
|
||||
self.currentAvailability = currentAvailability
|
||||
self.requestType = requestType
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""Converts the object to a JSON string matching the Java structure"""
|
||||
return json.dumps(
|
||||
{
|
||||
"version": self.version,
|
||||
"serviceId": self.serviceId,
|
||||
"taskId": self.taskId,
|
||||
"maxAvailability": self.maxAvailability,
|
||||
"currentAvailability": self.currentAvailability,
|
||||
"requestType": self.requestType.value, # Using .value for Enum serialization
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
|
||||
|
||||
class BackpressureHandler:
|
||||
# Track the number of messages currently being processed
|
||||
current_parallel_executions = 0
|
||||
last_data_message_time = 0 # helps detect IDLE condition
|
||||
last_backpressure_event_time = (
|
||||
0 # helps prevent flooding the system with backpressure events
|
||||
)
|
||||
last_backpressure_event = ScalingRequestAlert.IDLE
|
||||
|
||||
parallel_workers = int(os.getenv("PARALLEL_WORKERS", default=0))
|
||||
swarm_service_id = os.getenv("SWARM_SERVICE_ID", default="clever-amqp-service")
|
||||
swarm_task_id = os.getenv("SWARM_TASK_ID", default="python-adapter")
|
||||
|
||||
def __init__(self, channel: AbstractRobustChannel, loop: AbstractEventLoop):
|
||||
self.channel = channel
|
||||
self.loop = loop
|
||||
self.exchange = None
|
||||
|
||||
def handle_backpressure(self, scaling_request: ScaleRequestV1):
|
||||
# Handle the backpressure logic here
|
||||
print(
|
||||
f"Handling backpressure for {scaling_request.serviceId} with request type {scaling_request.requestType}"
|
||||
)
|
||||
|
||||
async def handle_backpressure_overload_event(self):
|
||||
logging.warning("Warning: Capacity close to depleted!")
|
||||
# Send an Overload event to the Management service
|
||||
scaling_request: ScaleRequestV1 = ScaleRequestV1(
|
||||
self.swarm_service_id,
|
||||
self.swarm_task_id,
|
||||
self.parallel_workers,
|
||||
self.parallel_workers - self.current_parallel_executions,
|
||||
ScalingRequestAlert.OVERLOAD,
|
||||
)
|
||||
# Address the message to any management service instance capable of supporting BACKPRESSURE request
|
||||
await self.publish_backpressure_request(scaling_request)
|
||||
# update / reset time-window so that the OVERLOAD is not sent too often
|
||||
self.last_data_message_time = time.time()
|
||||
|
||||
def _handle_backpressure_idle_event(self):
|
||||
logging.warning("Warning: Capacity idle!")
|
||||
# Send an Idle event to the Management service
|
||||
scaling_request: ScaleRequestV1 = ScaleRequestV1(
|
||||
self.swarm_service_id,
|
||||
self.swarm_task_id,
|
||||
self.parallel_workers,
|
||||
self.parallel_workers - self.current_parallel_executions,
|
||||
ScalingRequestAlert.IDLE,
|
||||
)
|
||||
# Address the message to any adapter capable of supporting BACKPRESSURE request
|
||||
self.publish_backpressure_request(scaling_request)
|
||||
# update / reset time-window so that the IDLE is not sent too often
|
||||
self.last_data_message_time = time.time()
|
||||
|
||||
async def publish_backpressure_request(self, scaling_request: ScaleRequestV1):
|
||||
# Publish the backpressure request to the management service
|
||||
if not self.exchange:
|
||||
|
||||
async def _wrap_rabbit_mq_api_init(channel):
|
||||
_exchange = await channel.declare_exchange(
|
||||
name="cleverthis.clevermicro.management", type=ExchangeType.fanout
|
||||
)
|
||||
return _exchange
|
||||
|
||||
_future = asyncio.run_coroutine_threadsafe(
|
||||
_wrap_rabbit_mq_api_init(self.channel), self.loop
|
||||
)
|
||||
while not _future.done():
|
||||
await asyncio.sleep(0.05)
|
||||
self.exchange = _future.result()
|
||||
|
||||
if self.exchange:
|
||||
|
||||
async def _wrap_rabbit_mq_api():
|
||||
if not self.channel.is_closed:
|
||||
binary_content: bytes = scaling_request.to_json().encode("utf-8")
|
||||
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 self.exchange.publish(message=pika_message, routing_key="")
|
||||
logging.info(
|
||||
" [x] Service Message Sent to %s, msg: %s",
|
||||
self.exchange.name,
|
||||
str(binary_content),
|
||||
)
|
||||
return True
|
||||
return False
|
||||
|
||||
_future = asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop)
|
||||
while not _future.done():
|
||||
await asyncio.sleep(0.05)
|
||||
@@ -7,17 +7,12 @@ import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import traceback
|
||||
from asyncio import AbstractEventLoop
|
||||
from typing import Dict, List, Tuple, Callable, Coroutine, Any
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
from xml.etree import ElementTree
|
||||
from typing import Dict, List, Callable, Coroutine, Any
|
||||
|
||||
from aio_pika.abc import (
|
||||
AbstractRobustChannel,
|
||||
AbstractRobustConnection,
|
||||
AbstractRobustExchange,
|
||||
AbstractExchange,
|
||||
)
|
||||
from aiohttp import ClientSession, ClientResponse
|
||||
@@ -25,18 +20,14 @@ from dataclasses import dataclass
|
||||
from fastapi import HTTPException
|
||||
from starlette.responses import FileResponse
|
||||
|
||||
from amqp.adapter.data_parser import AMQDataParser
|
||||
from amqp.adapter.data_response_factory import DataResponseFactory
|
||||
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
|
||||
from amqp.adapter.logging_utils import logging_error
|
||||
from amqp.config.amq_configuration import AMQConfiguration
|
||||
from amqp.model.model import DataMessage, DataResponse
|
||||
from amqp.adapter.multipart_parser import CleverMultiPartParser, process_form_data
|
||||
from amqp.adapter.file_uploader import publish_file_to_rabbitmq
|
||||
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
|
||||
|
||||
from IPython.core import debugger
|
||||
|
||||
debug = debugger.Pdb().set_trace
|
||||
|
||||
|
||||
@dataclass
|
||||
class AMQMessage(DataMessage):
|
||||
@@ -62,93 +53,10 @@ class AMQMessage(DataMessage):
|
||||
)
|
||||
|
||||
|
||||
def parse_get_url(url: str) -> Tuple[str, Dict[str, str]]:
|
||||
"""
|
||||
Parses an HTTP GET URL and returns (context_path, query_params_dict)
|
||||
|
||||
Args:
|
||||
url: The URL string to parse (e.g., "http://example.com/path?foo=1&bar=2&foo=3")
|
||||
|
||||
Returns:
|
||||
tuple: (context_path, query_params_dict)
|
||||
- context_path: The path part of the URL (e.g., "/path")
|
||||
- query_params_dict: Query parameters where values are str for single value or list of str for multivalue
|
||||
"""
|
||||
_parsed = urlparse(url)
|
||||
_query_params = parse_qs(_parsed.query)
|
||||
# Convert values from lists to single values when there's only one value
|
||||
_simplified_params = {
|
||||
k: v[0] if len(v) == 1 else v for k, v in _query_params.items()
|
||||
}
|
||||
return _parsed.path, _simplified_params
|
||||
|
||||
|
||||
def parse_request_body(message: DataMessage):
|
||||
"""
|
||||
Parses the HTTP POST request body based on the MIME type.
|
||||
This is typically called fom the concrete CleverThis Service to extract parameters from the message body.
|
||||
Args:
|
||||
message (DataMessage): Contains raw request and MIME headers.
|
||||
|
||||
Returns:
|
||||
dict: Parsed data as a dictionary.
|
||||
"""
|
||||
_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:
|
||||
logging_error(f"Invalid JSON: {e}")
|
||||
raise ValueError(f"Invalid JSON: {e}")
|
||||
|
||||
# Parse URL-encoded form data
|
||||
elif content_type == "application/x-www-form-urlencoded":
|
||||
try:
|
||||
_form_data: dict = {
|
||||
k: v[0] if len(v) == 1 else v
|
||||
for k, v in parse_qs(message.body()).items()
|
||||
}
|
||||
if len(_form_data) == 0 and len(message.body()) > 0:
|
||||
_form_data = json.loads(message.body())
|
||||
if _form_data["files"] and _form_data["form"]:
|
||||
_form_data = process_form_data(_form_data)
|
||||
return _form_data
|
||||
except Exception as e:
|
||||
logging_error(f"Invalid URL-encoded form data: {e}")
|
||||
raise ValueError(f"Invalid URL-encoded form data: {e}")
|
||||
|
||||
# Parse multipart/form-data
|
||||
elif content_type.startswith("multipart/form-data"):
|
||||
try:
|
||||
parser: CleverMultiPartParser = CleverMultiPartParser(message)
|
||||
return parser.form_data
|
||||
except Exception as e:
|
||||
logging_error(f"Invalid multipart/form-data: {e}")
|
||||
raise ValueError(f"Invalid multipart/form-data: {e}")
|
||||
|
||||
# Parse plain text
|
||||
elif content_type == "text/plain":
|
||||
return {"text": message.body()}
|
||||
|
||||
# Parse XML
|
||||
elif content_type == "application/xml":
|
||||
try:
|
||||
_root = ElementTree.fromstring(message.body())
|
||||
return {elem.tag: elem.text for elem in _root}
|
||||
except ElementTree.ParseError as e:
|
||||
raise ValueError(f"Invalid XML: {e}")
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported Content-Type: {content_type}")
|
||||
|
||||
|
||||
async def _convert_file_response(
|
||||
obj: FileResponse,
|
||||
connection: AbstractRobustConnection,
|
||||
file_handler: FileHandler,
|
||||
loop: AbstractEventLoop | None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -187,7 +95,7 @@ async def _convert_file_response(
|
||||
) # allow coroutine to execute. calling _future.result() will block the entire thread and the coroutine will never execute
|
||||
_channel, _exchange, _routing_key = _future.result()
|
||||
# publish_file also needs RabbitMQ API, ensure the loop is passed as argument
|
||||
(_stream_name, _file_size) = await publish_file_to_rabbitmq(
|
||||
(_stream_name, _file_size) = await file_handler.publish_file(
|
||||
_exchange, obj.path, _routing_key, loop
|
||||
)
|
||||
# as above, but luckily this time we don't need to wait for the result
|
||||
@@ -214,106 +122,6 @@ def is_likely_json(text: str) -> bool:
|
||||
)
|
||||
|
||||
|
||||
async def call_cleverthis_api(
|
||||
message: AMQMessage,
|
||||
args: Any,
|
||||
api_method: Callable[[Any, Any], Coroutine],
|
||||
success_code: int,
|
||||
loop: AbstractEventLoop | None = None,
|
||||
) -> DataResponse:
|
||||
"""
|
||||
Handles POST and PUT requests, which typically contain possibly complex payload data, passed as a dictionary.
|
||||
Invokes the provided API method with the given arguments and authenticated user,
|
||||
and returns the response as DataResponse.
|
||||
:param message: DataMessage containing the details of the call and the payload
|
||||
:param args: map of the arguments to pass to the API method
|
||||
:param api_method: Callable - the method to invoke. Last parameter is the authenticated user
|
||||
:param success_code: HTTP code to be returned when operation suceeds
|
||||
:return: DataResponse class
|
||||
"""
|
||||
try:
|
||||
_verified_args = {}
|
||||
_sig = inspect.signature(api_method)
|
||||
_required_args = list(_sig.parameters.values())
|
||||
# _required_args = getattr(api_method, '__annotations__', {})
|
||||
for arg in _required_args:
|
||||
if arg.name in args:
|
||||
# convert to another type if arg is not a str
|
||||
if arg.annotation == int:
|
||||
_verified_args[arg.name] = int(args[arg.name])
|
||||
elif arg.annotation == float:
|
||||
_verified_args[arg.name] = float(args[arg.name])
|
||||
elif arg.annotation == bool:
|
||||
_verified_args[arg.name] = bool(args[arg.name])
|
||||
# TODO: more?
|
||||
else:
|
||||
_verified_args[arg.name] = args[arg.name]
|
||||
else:
|
||||
logging.info(
|
||||
f"{api_method.__name__}: Missing required argument: {arg.name}"
|
||||
)
|
||||
logging.info(
|
||||
f" [*] INVOKE ENDPOINT: {api_method.__name__}, ARGS: {', '.join(f'{k}={v}' for k, v in _verified_args.items())}"
|
||||
)
|
||||
_result: Any = await api_method(**_verified_args)
|
||||
logging.info(
|
||||
f" [*] ENDPOINT RESULT: {api_method.__name__}, RESULT: {str(_result)}"
|
||||
)
|
||||
|
||||
if isinstance(_result, FileResponse):
|
||||
_json_result, _content_type = (
|
||||
await _convert_file_response(_result, message.connection, loop),
|
||||
"application/octet-stream",
|
||||
)
|
||||
else:
|
||||
_json_result = _result
|
||||
_content_type = "application/json"
|
||||
|
||||
_binary_result = (
|
||||
_json_result.encode("utf-8")
|
||||
if hasattr(_json_result, "encode")
|
||||
else (
|
||||
_json_result.body
|
||||
if hasattr(_json_result, "body")
|
||||
else (
|
||||
_json_result.read()
|
||||
if hasattr(_json_result, "read")
|
||||
else _json_result
|
||||
)
|
||||
)
|
||||
)
|
||||
logging.info(
|
||||
f" [x] ALL DONE in call_cleverthis_api id:{message.id}, resp code:{success_code}"
|
||||
)
|
||||
return DataResponseFactory.of(
|
||||
message.id, success_code, _content_type, _binary_result, message.trace_info
|
||||
)
|
||||
except HTTPException as http_error:
|
||||
_content = str(http_error.detail)
|
||||
if not is_likely_json(_content):
|
||||
_content = '{"error": "' + _content + '"}'
|
||||
logging_error(f"Service REST endpoint invocation failed: {http_error}")
|
||||
return DataResponseFactory.of(
|
||||
message.id,
|
||||
http_error.status_code,
|
||||
"application/json",
|
||||
_content.encode("utf-8"),
|
||||
message.trace_info,
|
||||
)
|
||||
except Exception as error:
|
||||
logging_error(f"Service endpoint invocation failed: {error}")
|
||||
_content = str(error)
|
||||
if not is_likely_json(_content):
|
||||
_content = '{"error": "' + _content + '"}'
|
||||
return DataResponseFactory.of(
|
||||
message.id,
|
||||
500,
|
||||
"application/json",
|
||||
_content.encode("utf-8"),
|
||||
message.trace_info,
|
||||
)
|
||||
|
||||
|
||||
# ============= CleverThisServiceAdapter =============
|
||||
class CleverThisServiceAdapter:
|
||||
"""
|
||||
@@ -336,13 +144,23 @@ class CleverThisServiceAdapter:
|
||||
self.service_host = self.amq_configuration.dispatch.service_host
|
||||
self.service_port = self.amq_configuration.dispatch.service_port
|
||||
self.loop: AbstractEventLoop | None = None
|
||||
self.require_authenticated_user = (
|
||||
eval(
|
||||
self.amq_configuration.amq_adapter.require_authenticated_user.capitalize()
|
||||
try:
|
||||
self.require_authenticated_user = (
|
||||
self.amq_configuration.amq_adapter.require_authenticated_user
|
||||
and isinstance(
|
||||
self.amq_configuration.amq_adapter.require_authenticated_user, bool
|
||||
)
|
||||
or eval(
|
||||
str(self.amq_configuration.amq_adapter.require_authenticated_user)
|
||||
.strip()
|
||||
.capitalize()
|
||||
or "True"
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
self.require_authenticated_user = (
|
||||
True # Enforce authenticated user as default
|
||||
)
|
||||
if self.amq_configuration.amq_adapter.require_authenticated_user.strip()
|
||||
else True
|
||||
)
|
||||
|
||||
def get_content_type(self, message_headers: Dict[str, List[str]]) -> str:
|
||||
"""
|
||||
@@ -384,7 +202,7 @@ class CleverThisServiceAdapter:
|
||||
if "Bearer" in _auth:
|
||||
try:
|
||||
token = _auth.split(" ")[1]
|
||||
logging.warn(f"Token: {token}")
|
||||
logging.warning(f"Token: {token}")
|
||||
_auth_user = await self.get_current_user(token)
|
||||
except Exception as error:
|
||||
logging_error(f"on_message: Error getting user: {error}")
|
||||
@@ -491,7 +309,7 @@ class CleverThisServiceAdapter:
|
||||
async def handle_possible_form(
|
||||
self, message: DataMessage, request_coroutine, auth_user: Any
|
||||
) -> DataResponse:
|
||||
_args = parse_request_body(message)
|
||||
_args = AMQDataParser.parse_request_body(message)
|
||||
return await request_coroutine
|
||||
|
||||
async def head(self, message: DataMessage, auth_user: Any) -> DataResponse:
|
||||
@@ -513,3 +331,112 @@ class CleverThisServiceAdapter:
|
||||
:return: DataResponse to be sent back to the API Gateway via AMQP
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
async def call_cleverthis_api(
|
||||
message: AMQMessage,
|
||||
args: Any,
|
||||
api_method: Callable[[Any, Any], Coroutine],
|
||||
success_code: int,
|
||||
loop: AbstractEventLoop | None = None,
|
||||
) -> DataResponse:
|
||||
"""
|
||||
Handles POST and PUT requests, which typically contain possibly complex payload data, passed as a dictionary.
|
||||
Invokes the provided API method with the given arguments and authenticated user,
|
||||
and returns the response as DataResponse.
|
||||
:param message: DataMessage containing the details of the call and the payload
|
||||
:param args: map of the arguments to pass to the API method
|
||||
:param api_method: Callable - the method to invoke. Last parameter is the authenticated user
|
||||
:param success_code: HTTP code to be returned when operation suceeds
|
||||
:return: DataResponse class
|
||||
"""
|
||||
try:
|
||||
_verified_args = {}
|
||||
_sig = inspect.signature(api_method)
|
||||
_required_args = list(_sig.parameters.values())
|
||||
# _required_args = getattr(api_method, '__annotations__', {})
|
||||
for arg in _required_args:
|
||||
if arg.name in args:
|
||||
# convert to another type if arg is not a str
|
||||
if arg.annotation == int:
|
||||
_verified_args[arg.name] = int(args[arg.name])
|
||||
elif arg.annotation == float:
|
||||
_verified_args[arg.name] = float(args[arg.name])
|
||||
elif arg.annotation == bool:
|
||||
_verified_args[arg.name] = bool(args[arg.name])
|
||||
# TODO: more?
|
||||
else:
|
||||
_verified_args[arg.name] = args[arg.name]
|
||||
else:
|
||||
logging.info(
|
||||
f"{api_method.__name__}: Missing required argument: {arg.name}"
|
||||
)
|
||||
logging.info(
|
||||
f" [*] INVOKE ENDPOINT: {api_method.__name__}, ARGS: {', '.join(f'{k}={v}' for k, v in _verified_args.items())}"
|
||||
)
|
||||
_result: Any = await api_method(**_verified_args)
|
||||
logging.info(
|
||||
f" [*] ENDPOINT RESULT: {api_method.__name__}, RESULT: {str(_result)}"
|
||||
)
|
||||
|
||||
if isinstance(_result, FileResponse):
|
||||
_json_result, _content_type = (
|
||||
await _convert_file_response(
|
||||
_result,
|
||||
message.connection,
|
||||
file_handler=StreamingFileHandler(),
|
||||
loop=loop,
|
||||
),
|
||||
"application/octet-stream",
|
||||
)
|
||||
else:
|
||||
_json_result = _result
|
||||
_content_type = "application/json"
|
||||
|
||||
_binary_result = (
|
||||
_json_result.encode("utf-8")
|
||||
if hasattr(_json_result, "encode")
|
||||
else (
|
||||
_json_result.body
|
||||
if hasattr(_json_result, "body")
|
||||
else (
|
||||
_json_result.read()
|
||||
if hasattr(_json_result, "read")
|
||||
else _json_result
|
||||
)
|
||||
)
|
||||
)
|
||||
logging.info(
|
||||
f" [x] ALL DONE in call_cleverthis_api id:{message.id}, resp code:{success_code}"
|
||||
)
|
||||
return DataResponseFactory.of(
|
||||
message.id,
|
||||
success_code,
|
||||
_content_type,
|
||||
_binary_result,
|
||||
message.trace_info,
|
||||
)
|
||||
except HTTPException as http_error:
|
||||
_content = str(http_error.detail)
|
||||
if not is_likely_json(_content):
|
||||
_content = '{"error": "' + _content + '"}'
|
||||
logging_error(f"Service REST endpoint invocation failed: {http_error}")
|
||||
return DataResponseFactory.of(
|
||||
message.id,
|
||||
http_error.status_code,
|
||||
"application/json",
|
||||
_content.encode("utf-8"),
|
||||
message.trace_info,
|
||||
)
|
||||
except Exception as error:
|
||||
logging_error(f"Service endpoint invocation failed: {error}")
|
||||
_content = str(error)
|
||||
if not is_likely_json(_content):
|
||||
_content = '{"error": "' + _content + '"}'
|
||||
return DataResponseFactory.of(
|
||||
message.id,
|
||||
500,
|
||||
"application/json",
|
||||
_content.encode("utf-8"),
|
||||
message.trace_info,
|
||||
)
|
||||
|
||||
@@ -9,6 +9,7 @@ from io import BytesIO
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, List
|
||||
|
||||
import amqp.adapter.file_handler
|
||||
from amqp.model.model import DataMessage
|
||||
from amqp.model.snowflake_id import SnowflakeId
|
||||
from amqp.adapter.serializer import (
|
||||
@@ -18,7 +19,7 @@ from amqp.adapter.serializer import (
|
||||
parse_map_string,
|
||||
)
|
||||
from amqp.router.utils import sanitize_as_rabbitmq_name
|
||||
from amqp.adapter.file_downloader import download_buffer
|
||||
from amqp.adapter.file_handler import StreamingFileHandler, FileHandler
|
||||
|
||||
|
||||
class DataMessageFactory:
|
||||
@@ -230,20 +231,27 @@ class DataMessageFactory:
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def from_stream(stream: bytes, rabbitmq_url: str, loop) -> DataMessage | None:
|
||||
async def from_stream(
|
||||
stream: bytes,
|
||||
rabbitmq_url: str,
|
||||
loop,
|
||||
file_handler: FileHandler = StreamingFileHandler(),
|
||||
) -> DataMessage | None:
|
||||
"""
|
||||
Builds the DataMessage from its serialized (byte array) form.
|
||||
:param stream: serialized message
|
||||
:param rabbitmq_url: RabbitMQ connection URL
|
||||
:param loop: event loop
|
||||
:param file_handler: file handler for downloading the message
|
||||
:return: DataMessage 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
|
||||
_message_data = json.loads(stream.decode())
|
||||
_req_info = _message_data.get("req-info", "")
|
||||
if _req_info:
|
||||
(body, eof) = await file_handler.download_buffer(
|
||||
loop=loop, rabbitmq_url=rabbitmq_url, queue_name=_req_info
|
||||
)
|
||||
if eof == 1:
|
||||
if eof == amqp.adapter.file_handler.END_OF_FILE:
|
||||
return DataMessageFactory.from_bytes(body)
|
||||
return None
|
||||
|
||||
@@ -0,0 +1,201 @@
|
||||
import io
|
||||
import json
|
||||
import python_multipart
|
||||
|
||||
from fastapi import UploadFile
|
||||
from python_multipart.multipart import Field, File
|
||||
from typing import Tuple, Dict
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
from xml.etree import ElementTree
|
||||
|
||||
from amqp.adapter.logging_utils import logging_error, logging_debug
|
||||
from amqp.model.model import DataMessage
|
||||
|
||||
|
||||
class MultipartFormDataParser:
|
||||
"""
|
||||
Parses multipart/form-data content and extracts form fields (key/value pairs) and files.
|
||||
Note: multipart message that contains file(s) to upload are converted to CleverMicro JSON-based streaming format,
|
||||
where files are transmitted separately via stream using dedicated queue, so the message content is no longer
|
||||
the HTTP multipart format, but the Content-Type is being preserved for compatibility with the original message.
|
||||
This class supports both formats (multipart and JSON) depending on the message content, because if no file is
|
||||
included, the message is transmitted in the unchanged multipart/form-data format.
|
||||
"""
|
||||
|
||||
def __init__(self, message: DataMessage):
|
||||
self.form_data: dict = {}
|
||||
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()
|
||||
}
|
||||
try:
|
||||
if (
|
||||
hasattr(message, "extra_data")
|
||||
and isinstance(message.extra_data, dict)
|
||||
and len(message.extra_data) > 0
|
||||
):
|
||||
try:
|
||||
self.form_data = json.loads(message.body())
|
||||
self.is_json = True
|
||||
if self.form_data["files"] and self.form_data["form"]:
|
||||
self.form_data = AMQDataParser.process_form_data(self.form_data)
|
||||
except Exception as jsonerror:
|
||||
logging_error(f"Parsing JSON: {jsonerror}")
|
||||
self.is_valid = False
|
||||
else:
|
||||
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
|
||||
if self.form_data["files"] and self.form_data["form"]:
|
||||
self.form_data = AMQDataParser.process_form_data(self.form_data)
|
||||
except Exception as jsonerror:
|
||||
logging_error(f"Parsing JSON: {jsonerror}")
|
||||
self.is_valid = False
|
||||
|
||||
def on_field(self, field: Field):
|
||||
logging_debug(str(field))
|
||||
self.form_data[field.field_name.decode("utf-8")] = field.value
|
||||
|
||||
def on_file(self, file: File):
|
||||
logging_debug(str(file))
|
||||
self.form_data[file.field_name.decode("utf-8")] = UploadFile(
|
||||
file.file_object, size=file.size, filename=file.file_name
|
||||
)
|
||||
|
||||
|
||||
class AMQDataParser:
|
||||
"""
|
||||
A class for parsing AMQ data messages. It provides methods to parse form data,
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def process_form_data(json_input) -> dict:
|
||||
"""
|
||||
Processes the input JSON to extract form data. The JSON has two top entries: 'files' and 'form'.
|
||||
This is part of the file stream implementation in AMQ adapter to stream large files in chunks.
|
||||
'files' contains information about the AMQ queue from which read the file, while 'form' contains
|
||||
additional non-file form values.
|
||||
This function combines the file and values at the top level in the returned dictionary.
|
||||
Args:
|
||||
json_input: JSON string or dict containing the form and files data
|
||||
|
||||
Returns:
|
||||
dict: Processed form data with files data overriding non-empty values
|
||||
"""
|
||||
# Load JSON if input is string
|
||||
if isinstance(json_input, str):
|
||||
data = json.loads(json_input)
|
||||
else:
|
||||
data = json_input
|
||||
|
||||
result = {}
|
||||
|
||||
# Process form data - take first element of each list
|
||||
for key, value in data["form"].items():
|
||||
if value: # Only process if list is not empty
|
||||
result[key] = value[0]
|
||||
|
||||
# Override with files data where available
|
||||
for key, value in data["files"].items():
|
||||
if value: # Only override if files data is non-empty
|
||||
result[key] = value[0]
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def parse_get_url(url: str) -> Tuple[str, Dict[str, str]]:
|
||||
"""
|
||||
Parses an HTTP GET URL and returns (context_path, query_params_dict)
|
||||
|
||||
Args:
|
||||
url: The URL string to parse (e.g., "http://example.com/path?foo=1&bar=2&foo=3")
|
||||
|
||||
Returns:
|
||||
tuple: (context_path, query_params_dict)
|
||||
- context_path: The path part of the URL (e.g., "/path")
|
||||
- query_params_dict: Query parameters where values are str for single value or list of str for multivalue
|
||||
"""
|
||||
_parsed = urlparse(url)
|
||||
_query_params = parse_qs(_parsed.query)
|
||||
# Convert values from lists to single values when there's only one value
|
||||
_simplified_params = {
|
||||
k: v[0] if len(v) == 1 else v for k, v in _query_params.items()
|
||||
}
|
||||
return _parsed.path, _simplified_params
|
||||
|
||||
@staticmethod
|
||||
def parse_request_body(message: DataMessage):
|
||||
"""
|
||||
Parses the HTTP POST request body based on the MIME type.
|
||||
This is typically called fom the concrete CleverThis Service to extract parameters from the message body.
|
||||
Args:
|
||||
message (DataMessage): Contains raw request and MIME headers.
|
||||
|
||||
Returns:
|
||||
dict: Parsed data as a dictionary.
|
||||
"""
|
||||
_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:
|
||||
logging_error(f"Invalid JSON: {e}")
|
||||
raise ValueError(f"Invalid JSON: {e}")
|
||||
|
||||
# Parse URL-encoded form data
|
||||
elif content_type == "application/x-www-form-urlencoded":
|
||||
try:
|
||||
_form_data: dict = {
|
||||
k: v[0] if len(v) == 1 else v
|
||||
for k, v in parse_qs(message.body()).items()
|
||||
}
|
||||
if len(_form_data) == 0 and len(message.body()) > 0:
|
||||
_form_data = json.loads(message.body())
|
||||
if _form_data["files"] and _form_data["form"]:
|
||||
_form_data = MultipartFormDataParser.process_form_data(
|
||||
_form_data
|
||||
)
|
||||
return _form_data
|
||||
except Exception as e:
|
||||
logging_error(f"Invalid URL-encoded form data: {e}")
|
||||
raise ValueError(f"Invalid URL-encoded form data: {e}")
|
||||
|
||||
# Parse multipart/form-data
|
||||
elif content_type.startswith("multipart/form-data"):
|
||||
try:
|
||||
parser: MultipartFormDataParser = MultipartFormDataParser(message)
|
||||
return parser.form_data
|
||||
except Exception as e:
|
||||
logging_error(f"Invalid multipart/form-data: {e}")
|
||||
raise ValueError(f"Invalid multipart/form-data: {e}")
|
||||
|
||||
# Parse plain text
|
||||
elif content_type == "text/plain":
|
||||
return {"text": message.body()}
|
||||
|
||||
# Parse XML
|
||||
elif content_type == "application/xml":
|
||||
try:
|
||||
_root = ElementTree.fromstring(message.body())
|
||||
return {elem.tag: elem.text for elem in _root}
|
||||
except ElementTree.ParseError as e:
|
||||
raise ValueError(f"Invalid XML: {e}")
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported Content-Type: {content_type}")
|
||||
@@ -1,232 +0,0 @@
|
||||
"""
|
||||
implements the file downloader for the amqp adapter, where file is sent in chunks over RabbitMQ
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from asyncio import Future
|
||||
from collections import defaultdict
|
||||
from typing import Any
|
||||
|
||||
from aio_pika import connect_robust
|
||||
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustChannel, AbstractQueue
|
||||
|
||||
from amqp.adapter.logging_utils import logging_error
|
||||
|
||||
# Global dictionary to track completed downloads
|
||||
download_locations = defaultdict(str)
|
||||
IN_PROGRESS = 0
|
||||
END_OF_FILE = 1
|
||||
CANCELLED = 2
|
||||
|
||||
|
||||
def ensure_directory_exists(filepath) -> str:
|
||||
"""
|
||||
Ensures that the directory for the given filepath exists, and handles file versioning.
|
||||
|
||||
Args:
|
||||
filepath: The path to the file.
|
||||
|
||||
Returns:
|
||||
The original filepath if the file does not exist, or a modified filepath
|
||||
pointing to the next available version of the file if it does.
|
||||
"""
|
||||
# 1. Ensure directory exists
|
||||
directory = os.path.dirname(filepath)
|
||||
if directory and not os.path.exists(directory):
|
||||
os.makedirs(directory)
|
||||
|
||||
# 2. Handle file versioning
|
||||
if not os.path.exists(filepath):
|
||||
return filepath # File doesn't exist, return original path
|
||||
|
||||
# File exists, find next available version
|
||||
base, ext = os.path.splitext(filepath)
|
||||
version = 1
|
||||
while True:
|
||||
new_filepath = f"{base}.{version}{ext}"
|
||||
if not os.path.exists(new_filepath):
|
||||
return new_filepath # Found a free version, return this path
|
||||
version += 1
|
||||
|
||||
|
||||
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)
|
||||
filepath = ensure_directory_exists(
|
||||
filepath
|
||||
) # if filename exists, create a new (next) version
|
||||
download_locations[file_info["key"]] = filepath
|
||||
filename = os.path.basename(
|
||||
filepath
|
||||
) # if a new version was created, ensure to use that new version
|
||||
file_info["filename"] = filename
|
||||
|
||||
logging.info(
|
||||
f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}"
|
||||
)
|
||||
|
||||
async def wrap_rabbit_mq_calls() -> tuple[int, int] | None:
|
||||
connection = await connect_robust(rabbitmq_url, loop=loop)
|
||||
try:
|
||||
_file_size = 0
|
||||
_eof = IN_PROGRESS
|
||||
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)
|
||||
_file_size = _file_size + message.body_size
|
||||
_eof = int(message.headers.get("eof", END_OF_FILE))
|
||||
if _eof != IN_PROGRESS:
|
||||
break
|
||||
except Exception as e:
|
||||
logging_error(f"File download failed: {e}")
|
||||
finally:
|
||||
await connection.close()
|
||||
|
||||
return _file_size, _eof
|
||||
|
||||
if queue_name:
|
||||
_future: Future = asyncio.run_coroutine_threadsafe(
|
||||
wrap_rabbit_mq_calls(), loop=loop
|
||||
)
|
||||
while not _future.done():
|
||||
await asyncio.sleep(0.05)
|
||||
_local_file_size, _local_eof = _future.result()
|
||||
return _local_file_size, _local_eof
|
||||
|
||||
return 0, CANCELLED
|
||||
|
||||
|
||||
async def download_buffer(loop, rabbitmq_url, queue_name) -> (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.
|
||||
"""
|
||||
logging.info(f"Downloading buffer from stream: {queue_name}")
|
||||
if queue_name:
|
||||
|
||||
async def wrap_rabbit_mq_calls() -> tuple[bytearray, int]:
|
||||
_connection = await connect_robust(rabbitmq_url, loop=loop)
|
||||
_eof = IN_PROGRESS
|
||||
_res = bytearray()
|
||||
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"Downloading buffer: {e}")
|
||||
|
||||
await _connection.close()
|
||||
return _res, _eof
|
||||
|
||||
_future: Future = asyncio.run_coroutine_threadsafe(
|
||||
wrap_rabbit_mq_calls(), loop=loop
|
||||
)
|
||||
while not _future.done():
|
||||
await asyncio.sleep(0.05)
|
||||
_local_res, _local_eof = _future.result()
|
||||
return _local_res, _local_eof
|
||||
return 0, CANCELLED
|
||||
|
||||
|
||||
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(
|
||||
" [%s] No files to download in this message.", message.delivery_tag
|
||||
)
|
||||
# await message.ack()
|
||||
return False
|
||||
|
||||
# Create a task for each file download
|
||||
_tasks = []
|
||||
for file_info in files_to_download:
|
||||
logging.info(
|
||||
f" [{message.delivery_tag}] about to 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(
|
||||
f" [{message.delivery_tag}] 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. d-tag: {message.delivery_tag}")
|
||||
# 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"Building DataMessage: {e}")
|
||||
asyncio.run_coroutine_threadsafe(
|
||||
message.ack(), loop=loop
|
||||
) # ACK -> ignore the message (do not redeliver)
|
||||
|
||||
return download_locations
|
||||
@@ -0,0 +1,418 @@
|
||||
"""
|
||||
implements the file downloader for the amqp adapter, where file is sent in chunks over RabbitMQ
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from asyncio import Future, AbstractEventLoop
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
from aio_pika import connect_robust, Message
|
||||
from aio_pika.abc import (
|
||||
AbstractIncomingMessage,
|
||||
AbstractRobustChannel,
|
||||
AbstractQueue,
|
||||
AbstractExchange,
|
||||
DeliveryMode,
|
||||
)
|
||||
|
||||
from amqp.adapter.logging_utils import logging_error, logging_debug
|
||||
|
||||
# Global dictionary to track completed downloads
|
||||
download_locations = defaultdict(str)
|
||||
IN_PROGRESS = 0
|
||||
END_OF_FILE = 1
|
||||
CANCELLED = 2
|
||||
|
||||
|
||||
class FileHandler:
|
||||
"""
|
||||
The abstract base class for file handling. This is mainly to make the code testable, as there are only 2
|
||||
options for pushing (uploading) a file(s) to a service: as HTTP multipart message delivered via RabbitMQ
|
||||
or via stream in RabbitMQ. Downloading a file is always done via RabbitMQ stream. Thus the method signiture
|
||||
is tighly coupled to the RabbitMQ API in v1. Only alternative implementation is a mock object in unittests.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def ensure_directory_exists(self, filepath) -> str:
|
||||
"""
|
||||
Ensures that the directory for the given filepath exists, and handles file versioning.
|
||||
|
||||
Args:
|
||||
filepath: The path to the file.
|
||||
|
||||
Returns:
|
||||
The original filepath if the file does not exist, or a modified filepath
|
||||
pointing to the next available version of the file if it does.
|
||||
"""
|
||||
# 1. Ensure directory exists
|
||||
directory = os.path.dirname(filepath)
|
||||
if directory and not os.path.exists(directory):
|
||||
os.makedirs(directory)
|
||||
|
||||
# 2. Handle file versioning
|
||||
if not os.path.exists(filepath):
|
||||
return filepath # File doesn't exist, return original path
|
||||
|
||||
# File exists, find next available version
|
||||
base, ext = os.path.splitext(filepath)
|
||||
version = 1
|
||||
while True:
|
||||
new_filepath = f"{base}.{version}{ext}"
|
||||
if not os.path.exists(new_filepath):
|
||||
return new_filepath # Found a free version, return this path
|
||||
version += 1
|
||||
|
||||
async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (bytearray, int):
|
||||
"""
|
||||
Asynchronously downloads a bytearray 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.
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement this method.")
|
||||
|
||||
async def download_file(
|
||||
self, 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. File may consist of multiple chunks, so the function will wait until all chunks are received.
|
||||
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.
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement this method.")
|
||||
|
||||
async def on_message(
|
||||
self,
|
||||
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.
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement this method.")
|
||||
|
||||
async def publish_file(
|
||||
self,
|
||||
exchange: AbstractExchange,
|
||||
file_path: Path,
|
||||
routing_key: str,
|
||||
loop: AbstractEventLoop,
|
||||
max_chunk_size: int = 2**20, # 1MB default chunk size
|
||||
content_type: str = "application/octet-stream", # Default content type
|
||||
delivery_mode: DeliveryMode = DeliveryMode.PERSISTENT,
|
||||
) -> tuple[str, int]:
|
||||
"""
|
||||
Reads a file from the file system and publishes its content to a RabbitMQ exchange
|
||||
in chunks. It declares a queue with a random name, binds it to the specified
|
||||
exchange, and uses the queue name as the routing key.
|
||||
|
||||
Args:
|
||||
exchange: The aio_pika RobustChannel to use for communication with RabbitMQ.
|
||||
file_path: The path to the file to publish.
|
||||
max_chunk_size: The maximum size of each chunk (in bytes).
|
||||
routing_key: The name of the RabbitMQ exchange routing key to publish to correct queue.
|
||||
loop: correct event loop
|
||||
max_chunk_size: chunk size to send
|
||||
content_type: The content type of the file data.
|
||||
delivery_mode: The delivery mode (e.g., PERSISTENT or TRANSIENT).
|
||||
|
||||
Returns:
|
||||
The name of the randomly generated queue where the file content was published.
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement this method.")
|
||||
|
||||
|
||||
# The StreamingFileHandler class is a subclass of FileHandler that implements the download_file method
|
||||
class StreamingFileHandler(FileHandler):
|
||||
"""
|
||||
Handles file downloads from RabbitMQ streams.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
async def download_file(
|
||||
self, 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)
|
||||
filepath = self.ensure_directory_exists(
|
||||
filepath
|
||||
) # if filename exists, create a new (next) version
|
||||
download_locations[file_info["key"]] = filepath
|
||||
filename = os.path.basename(
|
||||
filepath
|
||||
) # if a new version was created, ensure to use that new version
|
||||
file_info["filename"] = filename
|
||||
|
||||
logging.info(
|
||||
f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}"
|
||||
)
|
||||
|
||||
async def wrap_rabbit_mq_calls() -> tuple[int, int] | None:
|
||||
connection = await connect_robust(rabbitmq_url, loop=loop)
|
||||
try:
|
||||
_file_size = 0
|
||||
_eof = IN_PROGRESS
|
||||
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)
|
||||
_file_size = _file_size + message.body_size
|
||||
_eof = int(message.headers.get("eof", END_OF_FILE))
|
||||
if _eof != IN_PROGRESS:
|
||||
break
|
||||
except Exception as e:
|
||||
logging_error(f"File download failed: {e}")
|
||||
finally:
|
||||
await connection.close()
|
||||
|
||||
return _file_size, _eof
|
||||
|
||||
if queue_name:
|
||||
_future: Future = asyncio.run_coroutine_threadsafe(
|
||||
wrap_rabbit_mq_calls(), loop=loop
|
||||
)
|
||||
while not _future.done():
|
||||
await asyncio.sleep(0.05)
|
||||
_local_file_size, _local_eof = _future.result()
|
||||
return _local_file_size, _local_eof
|
||||
|
||||
return 0, CANCELLED
|
||||
|
||||
async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (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.
|
||||
"""
|
||||
logging.info(f"Downloading buffer from stream: {queue_name}")
|
||||
if queue_name:
|
||||
|
||||
async def wrap_rabbit_mq_calls() -> tuple[bytearray, int]:
|
||||
_connection = await connect_robust(rabbitmq_url, loop=loop)
|
||||
_eof = IN_PROGRESS
|
||||
_res = bytearray()
|
||||
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"Downloading buffer: {e}")
|
||||
|
||||
await _connection.close()
|
||||
return _res, _eof
|
||||
|
||||
_future: Future = asyncio.run_coroutine_threadsafe(
|
||||
wrap_rabbit_mq_calls(), loop=loop
|
||||
)
|
||||
while not _future.done():
|
||||
await asyncio.sleep(0.05)
|
||||
_local_res, _local_eof = _future.result()
|
||||
return _local_res, _local_eof
|
||||
return 0, CANCELLED
|
||||
|
||||
async def on_message(
|
||||
self,
|
||||
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.
|
||||
json_data: The JSON data from the message.
|
||||
rabbitmq_url: The RabbitMQ connection URL.
|
||||
loop: The asyncio event loop of the RabbitMQ API
|
||||
output_dir
|
||||
"""
|
||||
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(
|
||||
" [%s] No files to download in this message.", message.delivery_tag
|
||||
)
|
||||
# await message.ack()
|
||||
return defaultdict()
|
||||
|
||||
# Create a task for each file download
|
||||
_tasks = []
|
||||
for file_info in files_to_download:
|
||||
logging.info(
|
||||
f" [{message.delivery_tag}] about to create task for {file_info}"
|
||||
)
|
||||
_task = asyncio.create_task(
|
||||
self.download_file(
|
||||
loop, rabbitmq_url, file_info, _message_data, output_dir
|
||||
)
|
||||
)
|
||||
_tasks.append(_task)
|
||||
logging.info(
|
||||
f" [{message.delivery_tag}] 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. d-tag: {message.delivery_tag}")
|
||||
# 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"Building DataMessage: {e}")
|
||||
asyncio.run_coroutine_threadsafe(
|
||||
message.ack(), loop=loop
|
||||
) # ACK -> ignore the message (do not redeliver)
|
||||
|
||||
return download_locations
|
||||
|
||||
async def publish_file(
|
||||
self,
|
||||
exchange: AbstractExchange,
|
||||
file_path: Path,
|
||||
routing_key: str,
|
||||
loop: AbstractEventLoop,
|
||||
max_chunk_size: int = 2**20, # 1MB default chunk size
|
||||
content_type: str = "application/octet-stream", # Default content type
|
||||
delivery_mode: DeliveryMode = DeliveryMode.PERSISTENT,
|
||||
) -> tuple[str, int]:
|
||||
"""
|
||||
Reads a file from the file system and publishes its content to a RabbitMQ exchange
|
||||
in chunks. It declares a queue with a random name, binds it to the specified
|
||||
exchange, and uses the queue name as the routing key.
|
||||
|
||||
Args:
|
||||
exchange: The aio_pika RobustChannel to use for communication with RabbitMQ.
|
||||
file_path: The path to the file to publish.
|
||||
max_chunk_size: The maximum size of each chunk (in bytes).
|
||||
routing_key: The name of the RabbitMQ exchange routing key to publish to correct queue.
|
||||
loop: correct event loop
|
||||
max_chunk_size: chunk size to send
|
||||
content_type: The content type of the file data.
|
||||
delivery_mode: The delivery mode (e.g., PERSISTENT or TRANSIENT).
|
||||
|
||||
Returns:
|
||||
The name of the randomly generated queue where the file content was published.
|
||||
"""
|
||||
if not file_path.is_file():
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
_file_size = os.path.getsize(file_path)
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
_offset = 0
|
||||
while True:
|
||||
_chunk = f.read(max_chunk_size)
|
||||
if not _chunk:
|
||||
break # End of file
|
||||
|
||||
message = Message(
|
||||
body=_chunk,
|
||||
content_type=content_type,
|
||||
delivery_mode=delivery_mode,
|
||||
headers={
|
||||
"file_name": file_path.name,
|
||||
"pos": _offset,
|
||||
"total_size": _file_size,
|
||||
"chunk_size": len(_chunk),
|
||||
"eof": (
|
||||
END_OF_FILE
|
||||
if max_chunk_size + _offset >= _file_size
|
||||
else IN_PROGRESS
|
||||
),
|
||||
},
|
||||
)
|
||||
logging.debug(
|
||||
f"Going to Publish chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
|
||||
)
|
||||
_future = asyncio.run_coroutine_threadsafe(
|
||||
exchange.publish(message, routing_key=routing_key), loop
|
||||
)
|
||||
while not _future.done():
|
||||
# sleep allows coroutine to execute. must wait for completion to ensure all chunks written in right order
|
||||
await asyncio.sleep(0.02)
|
||||
logging.debug(
|
||||
f"Waiting for chunk {_offset}/{_file_size} future= {_future}"
|
||||
)
|
||||
|
||||
_offset += len(_chunk)
|
||||
logging.debug(
|
||||
f"Published chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
|
||||
)
|
||||
|
||||
logging_debug(
|
||||
f"File {file_path} published to RabbitMQ exchange {exchange.name}, queue {routing_key}"
|
||||
)
|
||||
return routing_key, _file_size
|
||||
@@ -1,92 +0,0 @@
|
||||
import logging
|
||||
from asyncio import AbstractEventLoop
|
||||
|
||||
import aio_pika
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from pathlib import Path
|
||||
from aio_pika.abc import AbstractRobustChannel, AbstractRobustExchange, AbstractExchange
|
||||
from aio_pika import Message, DeliveryMode
|
||||
|
||||
from amqp.adapter.file_downloader import IN_PROGRESS, END_OF_FILE
|
||||
from amqp.adapter.logging_utils import logging_debug
|
||||
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
|
||||
|
||||
|
||||
async def publish_file_to_rabbitmq(
|
||||
exchange: AbstractExchange,
|
||||
file_path: Path,
|
||||
routing_key: str,
|
||||
loop: AbstractEventLoop,
|
||||
max_chunk_size: int = 2**20, # 1MB default chunk size
|
||||
content_type: str = "application/octet-stream", # Default content type
|
||||
delivery_mode: DeliveryMode = DeliveryMode.PERSISTENT,
|
||||
) -> tuple[str, int]:
|
||||
"""
|
||||
Reads a file from the file system and publishes its content to a RabbitMQ exchange
|
||||
in chunks. It declares a queue with a random name, binds it to the specified
|
||||
exchange, and uses the queue name as the routing key.
|
||||
|
||||
Args:
|
||||
exchange: The aio_pika RobustChannel to use for communication with RabbitMQ.
|
||||
file_path: The path to the file to publish.
|
||||
max_chunk_size: The maximum size of each chunk (in bytes).
|
||||
routing_key: The name of the RabbitMQ exchange routing key to publish to correct queue.
|
||||
loop: correct event loop
|
||||
max_chunk_size: chunk size to send
|
||||
content_type: The content type of the file data.
|
||||
delivery_mode: The delivery mode (e.g., PERSISTENT or TRANSIENT).
|
||||
|
||||
Returns:
|
||||
The name of the randomly generated queue where the file content was published.
|
||||
"""
|
||||
if not file_path.is_file():
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
_file_size = os.path.getsize(file_path)
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
_offset = 0
|
||||
while True:
|
||||
_chunk = f.read(max_chunk_size)
|
||||
if not _chunk:
|
||||
break # End of file
|
||||
|
||||
message = Message(
|
||||
body=_chunk,
|
||||
content_type=content_type,
|
||||
delivery_mode=delivery_mode,
|
||||
headers={
|
||||
"file_name": file_path.name,
|
||||
"pos": _offset,
|
||||
"total_size": _file_size,
|
||||
"chunk_size": len(_chunk),
|
||||
"eof": (
|
||||
END_OF_FILE
|
||||
if max_chunk_size + _offset >= _file_size
|
||||
else IN_PROGRESS
|
||||
),
|
||||
},
|
||||
)
|
||||
logging.debug(
|
||||
f"Going to Publish chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
|
||||
)
|
||||
_future = asyncio.run_coroutine_threadsafe(
|
||||
exchange.publish(message, routing_key=routing_key), loop
|
||||
)
|
||||
while not _future.done():
|
||||
# sleep allows coroutine to execute. must wait for completion to ensure all chunks written in right order
|
||||
await asyncio.sleep(0.02)
|
||||
logging.debug(
|
||||
f"Waiting for chunk {_offset}/{_file_size} future= {_future}"
|
||||
)
|
||||
|
||||
_offset += len(_chunk)
|
||||
logging.debug(
|
||||
f"Published chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
|
||||
)
|
||||
|
||||
logging_debug(
|
||||
f"File {file_path} published to RabbitMQ exchange {exchange.name}, queue {routing_key}"
|
||||
)
|
||||
return routing_key, _file_size
|
||||
@@ -1,111 +0,0 @@
|
||||
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.adapter.logging_utils import logging_error, logging_debug
|
||||
from amqp.model.model import DataMessage
|
||||
|
||||
|
||||
def process_form_data(json_input) -> dict:
|
||||
"""
|
||||
Processes the input JSON to extract form data. The JSON has two top entries: 'files' and 'form'.
|
||||
This is part of the file stream implementation in AMQ adapter to stream large files in chunks.
|
||||
'files' contains information about the AMQ queue from which read the file, while 'form' contains
|
||||
additional non-file form values.
|
||||
This function combines the file and values at the top level in the returned dictionary.
|
||||
Args:
|
||||
json_input: JSON string or dict containing the form and files data
|
||||
|
||||
Returns:
|
||||
dict: Processed form data with files data overriding non-empty values
|
||||
"""
|
||||
# Load JSON if input is string
|
||||
if isinstance(json_input, str):
|
||||
data = json.loads(json_input)
|
||||
else:
|
||||
data = json_input
|
||||
|
||||
result = {}
|
||||
|
||||
# Process form data - take first element of each list
|
||||
for key, value in data["form"].items():
|
||||
if value: # Only process if list is not empty
|
||||
result[key] = value[0]
|
||||
|
||||
# Override with files data where available
|
||||
for key, value in data["files"].items():
|
||||
if value: # Only override if files data is non-empty
|
||||
result[key] = value[0]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class CleverMultiPartParser:
|
||||
def __init__(self, message: DataMessage):
|
||||
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:
|
||||
if (
|
||||
hasattr(message, "extra_data")
|
||||
and isinstance(message.extra_data, dict)
|
||||
and len(message.extra_data) > 0
|
||||
):
|
||||
try:
|
||||
self.form_data = json.loads(message.body())
|
||||
self.is_json = True
|
||||
if self.form_data["files"] and self.form_data["form"]:
|
||||
self.form_data = process_form_data(self.form_data)
|
||||
except Exception as jsonerror:
|
||||
logging_error(f"Parsing JSON: {jsonerror}")
|
||||
self.is_valid = False
|
||||
else:
|
||||
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
|
||||
if self.form_data["files"] and self.form_data["form"]:
|
||||
self.form_data = process_form_data(self.form_data)
|
||||
except Exception as jsonerror:
|
||||
logging_error(f"Parsing JSON: {jsonerror}")
|
||||
self.is_valid = False
|
||||
else:
|
||||
self.form_data = message.body()
|
||||
|
||||
def on_field(self, field: Field):
|
||||
logging_debug(field)
|
||||
self.form_data[field.field_name.decode("utf-8")] = field.value
|
||||
logging_debug(self.form_data)
|
||||
|
||||
def on_file(self, file: File):
|
||||
logging_debug(file)
|
||||
self.form_data[file.field_name.decode("utf-8")] = UploadFile(
|
||||
file.file_object, size=file.size, filename=file.file_name
|
||||
)
|
||||
logging_debug(self.form_data)
|
||||
@@ -0,0 +1,246 @@
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from aiohttp import ClientResponse
|
||||
|
||||
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage
|
||||
from amqp.config.amq_configuration import AMQConfiguration
|
||||
from amqp.model.model import DataMessage, DataResponse
|
||||
|
||||
|
||||
class TestCleverThisServiceAdapter(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# Mock AMQConfiguration
|
||||
self.mock_amq_config = AMQConfiguration("./application.properties.local")
|
||||
|
||||
# Mock ClientSession
|
||||
self.mock_session = AsyncMock()
|
||||
|
||||
# Create adapter instance
|
||||
self.adapter = CleverThisServiceAdapter(
|
||||
amq_configuration=self.mock_amq_config, session=self.mock_session
|
||||
)
|
||||
|
||||
def test_init(self):
|
||||
self.assertEqual(self.adapter.service_host, "localhost")
|
||||
self.assertEqual(self.adapter.service_port, 8080)
|
||||
self.assertFalse(self.adapter.require_authenticated_user)
|
||||
|
||||
def test_get_content_type(self):
|
||||
test_cases = [
|
||||
({"Content-Type": ["application/json"]}, "application/json"),
|
||||
({"content-type": ["text/xml"]}, "text/xml"),
|
||||
({"CONTENT-TYPE": ["text/plain; charset=utf-8"]}, "text/plain"),
|
||||
({}, "application/json"),
|
||||
({"Other-Header": ["value"]}, "application/json"),
|
||||
]
|
||||
|
||||
for headers, expected in test_cases:
|
||||
with self.subTest(headers=headers):
|
||||
result = self.adapter.get_content_type(headers)
|
||||
self.assertEqual(result, expected)
|
||||
|
||||
async def test_get_current_user(self):
|
||||
# Base implementation should return None
|
||||
result = await self.adapter.get_current_user("test_token")
|
||||
self.assertIsNone(result)
|
||||
|
||||
@patch.object(CleverThisServiceAdapter, "get_current_user")
|
||||
async def test_on_message_unauthenticated(self, mock_get_user):
|
||||
mock_message = MagicMock(spec=AMQMessage)
|
||||
mock_message.headers = {}
|
||||
mock_message.path = "/some/path"
|
||||
mock_message.method = "GET"
|
||||
mock_message.id = "123"
|
||||
mock_message.trace_info = {}
|
||||
|
||||
# Test when authentication is required but not provided
|
||||
response = await self.adapter.on_message(mock_message)
|
||||
self.assertEqual(response.response_code, 401)
|
||||
mock_get_user.assert_not_called()
|
||||
|
||||
@patch.object(CleverThisServiceAdapter, "get_current_user")
|
||||
@patch.object(CleverThisServiceAdapter, "get")
|
||||
async def test_on_message_authenticated(self, mock_get, mock_get_user):
|
||||
mock_message = MagicMock(spec=AMQMessage)
|
||||
mock_message.headers = {"Authorization": ["Bearer valid_token"]}
|
||||
mock_message.path = "/some/path"
|
||||
mock_message.method = "GET"
|
||||
mock_message.id = "123"
|
||||
mock_message.trace_info = {}
|
||||
|
||||
mock_user = MagicMock()
|
||||
mock_get_user.return_value = mock_user
|
||||
|
||||
mock_response = MagicMock(spec=DataResponse)
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
response = await self.adapter.on_message(mock_message)
|
||||
|
||||
self.assertEqual(response, mock_response)
|
||||
mock_get_user.assert_called_once_with("valid_token")
|
||||
mock_get.assert_called_once_with(mock_message, mock_user)
|
||||
|
||||
@patch.object(CleverThisServiceAdapter, "get_current_user")
|
||||
async def test_on_message_login_path(self, mock_get_user):
|
||||
mock_message = MagicMock(spec=AMQMessage)
|
||||
mock_message.headers = {}
|
||||
mock_message.path = "/login"
|
||||
mock_message.method = "POST"
|
||||
mock_message.id = "123"
|
||||
mock_message.trace_info = {}
|
||||
|
||||
mock_response = MagicMock(spec=DataResponse)
|
||||
with patch.object(
|
||||
self.adapter, "post", return_value=mock_response
|
||||
) as mock_post:
|
||||
response = await self.adapter.on_message(mock_message)
|
||||
self.assertEqual(response, mock_response)
|
||||
mock_get_user.assert_not_called()
|
||||
mock_post.assert_called_once()
|
||||
|
||||
@patch.object(CleverThisServiceAdapter, "get_current_user")
|
||||
async def test_on_message_all_methods(self, mock_get_user):
|
||||
mock_user = MagicMock()
|
||||
mock_get_user.return_value = mock_user
|
||||
|
||||
mock_response = MagicMock(spec=DataResponse)
|
||||
|
||||
methods = ["GET", "PUT", "POST", "DELETE", "HEAD"]
|
||||
for method in methods:
|
||||
with self.subTest(method=method):
|
||||
mock_message = MagicMock(spec=AMQMessage)
|
||||
mock_message.headers = {"Authorization": ["Bearer valid_token"]}
|
||||
mock_message.path = "/some/path"
|
||||
mock_message.method = method
|
||||
mock_message.id = "123"
|
||||
mock_message.trace_info = {}
|
||||
|
||||
with patch.object(
|
||||
self.adapter, method.lower(), return_value=mock_response
|
||||
) as mock_method:
|
||||
response = await self.adapter.on_message(mock_message)
|
||||
self.assertEqual(response, mock_response)
|
||||
mock_method.assert_called_once()
|
||||
|
||||
async def test_get_with_session(self):
|
||||
mock_message = MagicMock(spec=DataMessage)
|
||||
mock_message.path = "/api/test"
|
||||
mock_message.headers = {"X-Test": "value"}
|
||||
mock_message.trace_info = {"trace-id": "123"}
|
||||
mock_message.id = "msg123"
|
||||
|
||||
mock_response = MagicMock(spec=ClientResponse)
|
||||
mock_response.status = 200
|
||||
mock_response.content_type = "application/json"
|
||||
mock_response.read = AsyncMock(return_value=b'{"result": "success"}')
|
||||
|
||||
self.mock_session.get.return_value = mock_response
|
||||
|
||||
response = await self.adapter.get(mock_message, None)
|
||||
|
||||
self.assertEqual(response.response_code, 200)
|
||||
self.mock_session.get.assert_called_once_with(
|
||||
"http://testhost:8080/api/test",
|
||||
headers={"X-Test": "value", "trace-id": "123"},
|
||||
)
|
||||
|
||||
async def test_get_without_session(self):
|
||||
self.adapter.session = None
|
||||
|
||||
mock_message = MagicMock(spec=DataMessage)
|
||||
mock_message.path = "/api/test"
|
||||
mock_message.id = "msg123"
|
||||
mock_message.trace_info = {}
|
||||
|
||||
response = await self.adapter.get(mock_message, None)
|
||||
|
||||
self.assertEqual(response.response_code, 404)
|
||||
|
||||
async def test_put_with_form_handling(self):
|
||||
mock_message = MagicMock(spec=DataMessage)
|
||||
mock_message.path = "/api/test"
|
||||
mock_message.headers = {}
|
||||
mock_message.trace_info = {}
|
||||
mock_message.id = "msg123"
|
||||
mock_message.body = b"test=value"
|
||||
|
||||
mock_response = MagicMock(spec=ClientResponse)
|
||||
mock_response.status = 200
|
||||
mock_response.content_type = "application/json"
|
||||
mock_response.read = AsyncMock(return_value=b'{"result": "success"}')
|
||||
|
||||
self.mock_session.put.return_value = mock_response
|
||||
|
||||
with patch.object(
|
||||
self.adapter,
|
||||
"handle_possible_form",
|
||||
wraps=self.adapter.handle_possible_form,
|
||||
) as mock_handler:
|
||||
response = await self.adapter.put(mock_message, None)
|
||||
|
||||
self.assertEqual(response.response_code, 200)
|
||||
mock_handler.assert_called_once()
|
||||
|
||||
async def test_post_with_form_handling(self):
|
||||
mock_message = MagicMock(spec=DataMessage)
|
||||
mock_message.path = "/api/test"
|
||||
mock_message.headers = {}
|
||||
mock_message.trace_info = {}
|
||||
mock_message.id = "msg123"
|
||||
mock_message.body = b"test=value"
|
||||
|
||||
mock_response = MagicMock(spec=ClientResponse)
|
||||
mock_response.status = 201
|
||||
mock_response.content_type = "application/json"
|
||||
mock_response.read = AsyncMock(return_value=b'{"result": "created"}')
|
||||
|
||||
self.mock_session.post.return_value = mock_response
|
||||
|
||||
with patch.object(
|
||||
self.adapter,
|
||||
"handle_possible_form",
|
||||
wraps=self.adapter.handle_possible_form,
|
||||
) as mock_handler:
|
||||
response = await self.adapter.post(mock_message, None)
|
||||
|
||||
self.assertEqual(response.response_code, 201)
|
||||
mock_handler.assert_called_once()
|
||||
|
||||
async def test_head_not_implemented(self):
|
||||
mock_message = MagicMock(spec=DataMessage)
|
||||
mock_message.id = "msg123"
|
||||
|
||||
with self.assertRaises(NotImplementedError):
|
||||
await self.adapter.head(mock_message, None)
|
||||
|
||||
async def test_delete_not_implemented(self):
|
||||
mock_message = MagicMock(spec=DataMessage)
|
||||
mock_message.id = "msg123"
|
||||
|
||||
# Assuming delete is not implemented in base class
|
||||
with self.assertRaises(NotImplementedError):
|
||||
await self.adapter.delete(mock_message, None)
|
||||
|
||||
def test_require_authenticated_user_config(self):
|
||||
# Test different config values for require_authenticated_user
|
||||
test_cases = [
|
||||
("True", True),
|
||||
("False", False),
|
||||
("", True), # Default when empty
|
||||
("invalid", True), # Default when invalid
|
||||
]
|
||||
|
||||
for config_value, expected in test_cases:
|
||||
with self.subTest(config_value=config_value):
|
||||
mock_config = AMQConfiguration("./application.properties.local")
|
||||
mock_config.dispatch.service_host = "testhost"
|
||||
mock_config.dispatch.service_port = 8080
|
||||
mock_config.amq_adapter.require_authenticated_user = config_value
|
||||
|
||||
adapter = CleverThisServiceAdapter(mock_config)
|
||||
self.assertEqual(adapter.require_authenticated_user, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -2,12 +2,10 @@ import os
|
||||
import tempfile
|
||||
from unittest import TestCase
|
||||
|
||||
import pytest
|
||||
|
||||
from amqp.adapter.file_downloader import ensure_directory_exists
|
||||
from amqp.adapter.file_handler import StreamingFileHandler
|
||||
|
||||
|
||||
class TestFileDownloader(TestCase):
|
||||
class TestFileHandler(TestCase):
|
||||
def test_ensure_directory_exists(self):
|
||||
"""
|
||||
Test function for ensure_directory_exists. Creates a temporary directory
|
||||
@@ -15,10 +13,11 @@ class TestFileDownloader(TestCase):
|
||||
"""
|
||||
|
||||
# Create a temporary directory
|
||||
file_handler = StreamingFileHandler()
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Test case 1: File does not exist
|
||||
filepath1 = os.path.join(temp_dir, "test1.txt")
|
||||
result1 = ensure_directory_exists(filepath1)
|
||||
result1 = file_handler.ensure_directory_exists(filepath1)
|
||||
assert result1 == filepath1, f"Test Case 1 Failed: {result1} != {filepath1}"
|
||||
# Create the file
|
||||
with open(filepath1, "w") as f:
|
||||
@@ -26,11 +25,11 @@ class TestFileDownloader(TestCase):
|
||||
|
||||
# Test case 2: File exists, should create versioned file
|
||||
filepath2 = os.path.join(temp_dir, "test2.txt")
|
||||
result2 = ensure_directory_exists(filepath2)
|
||||
result2 = file_handler.ensure_directory_exists(filepath2)
|
||||
assert result2 == filepath2, f"Test Case 2 Failed: {result2} != {filepath2}"
|
||||
with open(filepath2, "w") as f:
|
||||
f.write("test")
|
||||
result2_v1 = ensure_directory_exists(filepath2)
|
||||
result2_v1 = file_handler.ensure_directory_exists(filepath2)
|
||||
assert result2_v1 == os.path.join(
|
||||
temp_dir, "test2.1.txt"
|
||||
), f"Test Case 2 Version 1 Failed: {result2_v1} != {os.path.join(temp_dir, 'test2.1.txt')}"
|
||||
@@ -43,7 +42,7 @@ class TestFileDownloader(TestCase):
|
||||
with open(os.path.join(temp_dir, "test3.2.txt"), "w") as f:
|
||||
f.write("test")
|
||||
filepath3 = os.path.join(temp_dir, "test3.txt")
|
||||
result3_v3 = ensure_directory_exists(filepath3)
|
||||
result3_v3 = file_handler.ensure_directory_exists(filepath3)
|
||||
assert result3_v3 == os.path.join(
|
||||
temp_dir, "test3.3.txt"
|
||||
), f"Test Case 3 Version 3 Failed: {result3_v3} != {os.path.join(temp_dir, 'test3.3.txt')}"
|
||||
@@ -1,7 +1,5 @@
|
||||
from unittest import TestCase
|
||||
|
||||
import pytest
|
||||
|
||||
from amqp.adapter.pydantic_serializer import parse_url_query
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user