Files
Stanislav Hejny 7fd1c62596
Unit test coverage / pytest (push) Failing after 1m17s
/ build-and-push (push) Successful in 1m29s
Feat #58 - fix backpressure logic and update tests
2025-06-28 10:20:55 +02:00

363 lines
10 KiB
Python

import base64
import json
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List
from aio_pika.abc import AbstractRobustConnection
from amqp.model.service_message_type import ServiceMessageType
from amqp.model.snowflake_id import SnowflakeId
from amqp.router.utils import sanitize_as_rabbitmq_name
RS = ":"
"""
Define the structure of messages used in Clever Micro
"""
class DataMessageMagicByte(Enum):
HTTP_REQUEST = "A"
RPC_REQUEST = "R"
@dataclass
class CleverMicroMessage:
DEFAULT_CONTENT_TYPE = ["application/octet-stream"]
"""
The common structure of a message that transfers the user request to the Service and back. All data related messages
conform to this base structure:
magic - identifies the message type (REST, RPC, DataResponse)
id - unique message ID
m_field - method name (for REST style calls) or service name (for RPC style calls) or response code for DataResponse
d_field - domain name (for REST style calls) or module name (for RPC style calls) or error code for DataResponse
p_field - path (for REST style calls) or method name (for RPC style calls) or error cause for DataResponse
headers - headers of the message (e.g. Content-Type, Accept, etc.)
trace_info - trace information (e.g. trace ID, span ID, etc. for Jaeger trace monitor)
base64body - the body of the message encoded in base64
"""
def __init__(
self,
magic: str,
id: SnowflakeId,
m_field: str,
d_field: str,
p_field: str,
headers: Dict[str, List[str]],
trace_info: Dict[str, str],
base64body: bytes,
):
self._magic = magic
self._id: SnowflakeId = id
self._m_field = m_field
self._d_field = d_field
self._p_field = p_field
self._headers: Dict[str, List[str]] = headers
self._trace_info: Dict[str, str] = trace_info
self._base64body: bytes = base64body
def magic(self):
return self._magic
def id(self):
return self._id
def m_field(self) -> str:
return self._m_field
def d_field(self) -> str:
return self._d_field
def p_field(self) -> str:
return self._p_field
def headers(self) -> Dict[str, List[str]]:
return self._headers
def trace_info(self) -> Dict[str, str]:
return self._trace_info
def base64body(self) -> bytes:
return self._base64body
def body(self) -> bytes:
return base64.b64decode(self.base64body())
def content_type(self) -> str:
return self.headers().get("Content-Type", self.DEFAULT_CONTENT_TYPE)[0]
def return_address(self) -> str:
return self.headers().get("Return-Address", "")[0]
class DataMessage(CleverMicroMessage):
"""
The class represents a structure of a message that transfers the user request to the Service.
It may also be used to make a call/request between services.
The response to this call is in format of DataResponse class.
See DataMessageFactory for code to instantiate, serialize and deserialize the class
"""
def __init__(
self,
magic: str,
id: SnowflakeId,
method: str,
domain: str,
path: str,
headers: Dict[str, List[str]],
trace_info: Dict[str, str],
base64body: bytes,
):
super().__init__(magic, id, method, domain, path, headers, trace_info, base64body)
def method(self) -> str:
return self.m_field()
def domain(self) -> str:
return self.d_field()
def path(self) -> str:
return self.p_field()
def routing_key(self) -> str:
return sanitize_as_rabbitmq_name(f"{self.d_field()}.{self.p_field()}")
class MultipartDataMessage(DataMessage):
def __init__(self, message: DataMessage, extra_data: dict):
super().__init__(
message.magic(),
message.id(),
message.method(),
message.domain(),
message.path(),
message.headers(),
message.trace_info(),
message.base64body(),
)
self.extra_data = extra_data
@dataclass
class AMQMessage(DataMessage):
"""
Expand the core DataMessage type to include the channel, which is required to create a dedicated queue
in case the response should initiate file download.
"""
_connection: AbstractRobustConnection = field(repr=False) # Don't show in repr
def __init__(self, data_message: DataMessage, connection: AbstractRobustConnection):
# Initialize the parent DataMessage with all properties from the input data_message
super().__init__(
magic=data_message.magic(),
id=data_message.id(),
method=data_message.method(),
domain=data_message.domain(),
path=data_message.path(),
headers=data_message.headers(),
trace_info=data_message.trace_info(),
base64body=data_message.base64body(),
)
self._connection = connection
self.extra_data = data_message.extra_data if hasattr(data_message, "extra_data") else {}
@property
def connection(self) -> AbstractRobustConnection:
"""Get the AMQP connection associated with this message"""
return self._connection
class DataResponse(CleverMicroMessage):
"""
The response message to REST & RPC style calls (note: the request is made with DataMessage)
See DataResponseFactory for code to instantiate, serialize and deserialize the class
"""
def __init__(
self,
magic: str,
id: SnowflakeId,
response_code: str,
content_type: str,
error: str,
error_cause: str,
headers: Dict[str, List[str]],
trace_info: Dict[str, str],
base64body: bytes,
):
if headers is None:
headers = {}
headers["Content-Type"] = [content_type]
super().__init__(
magic,
id,
response_code,
error,
error_cause,
headers,
trace_info,
base64body,
)
def response_code(self) -> str:
return self.m_field()
def error(self) -> str:
return self.d_field()
def error_cause(self) -> str:
return self.p_field()
@dataclass(frozen=True)
class AMQRoute:
"""
Configuration item representing one AMQRoute. Adapter uses this to declare the Exchange
and/or queue and bind the key as routing expression.
:param component_name: a service/component name that this route points to (owner of the route)
:param exchange: RabbitMQ exchange name
:param queue: RabbitMQ queue name
:param key: RabbitMQ TOPIC Exchange routing key
:param timeout: Timeout in milliseconds, how long to wait for reply, or not waiting if less then 1
:param valid_until: optional timestamp (millis from epoch) after which the route is ignored
"""
component_name: str
exchange: str
queue: str
key: str
timeout: int
valid_until: int
FORMAT = "Component-Name:Exchange-Name:Queue-Name:Routing-Key:Timeout:Valid-Until"
def __str__(self):
return (
f"{self.component_name}{RS}{self.exchange}{RS}{self.queue}{RS}{self.key}"
f"{RS}{self.timeout}{RS}{self.valid_until}"
)
def __eq__(self, other):
if isinstance(other, AMQRoute):
return (
self.component_name == other.component_name
and self.queue == other.queue
and self.key == other.key
)
return False
@property
def primary_key(self):
"""
Build the unique identifier of the route
:return primary key
"""
return f"{self.component_name}{RS}{self.queue}{RS}{self.key}"
def as_string(self) -> str:
return self.__str__()
# Ensure backward compatibility
AMQResponse = DataResponse
@dataclass(frozen=True)
class ServiceMessage:
"""
The CleverMicro service message that is used to communicate the internal state of the AMQ channel.
In v0.2 it supports service discovery / route registration and Backpressure status notification.
"""
id: SnowflakeId = None
payload_type: int = 0
message_type: ServiceMessageType = ServiceMessageType.UNKNOWN
recipient_name: str = None
reply_to: str = None
trace_info: Dict[str, str] = field(default_factory=dict)
message: str = None
def is_valid(self) -> bool:
return self.message_type is not None
class ScalingRequestAlert(Enum):
OVERLOAD = "OVERLOAD"
IDLE = "IDLE"
UPDATE = "UPDATE"
SCALE_UP = "SCALE_UP"
SCALE_DOWN = "SCALE_DOWN"
@dataclass(frozen=True)
class ScalingRequest:
service_id: str
task_id: str
max_availability: int
current_availability: int
request_type: ScalingRequestAlert
version: int = 1
def to_json(self) -> str:
"""Converts the object to a JSON string matching the Java structure"""
return json.dumps(
{
"version": self.version,
"serviceId": self.service_id,
"taskId": self.task_id,
"maxAvailability": self.max_availability,
"currentAvailability": self.current_availability,
"requestType": self.request_type.value, # Using .value for Enum serialization
},
indent=2,
)
@classmethod
def from_json(cls, json_str: str) -> "ScalingRequest":
"""Creates a ScalingRequest from a JSON string"""
data = json.loads(json_str)
return cls(
service_id=data["serviceId"],
task_id=data["taskId"],
max_availability=data["maxAvailability"],
current_availability=data["currentAvailability"],
request_type=ScalingRequestAlert(data["requestType"]), # Convert string to Enum
)
class AMQErrorMessage:
"""
The error response message to RPC style call (note: the request is made with DataMessage)
"""
def __init__(
self,
response_code: int,
error: str,
detail: str,
content_type: str = "application/json",
):
self.response_code = response_code
self.error = error
self.detail = detail
self.content_type = content_type
def serialize(self) -> bytes:
return json.dumps(
{
"response_code": self.response_code,
"error": self.error,
"detail": self.detail,
}
).encode("utf-8")