Files
amq-adapter-python/amqp/model/model.py
T
Stanislav Hejny 5bb6597b57
Unit test coverage / pytest (push) Failing after 1m1s
Unit test coverage / pytest (pull_request) Failing after 1m18s
#40 - ensure the JSON format is handled correctly at all levels
2025-05-07 13:05:04 +01:00

335 lines
8.7 KiB
Python

import base64
import json
from enum import Enum
from typing import List, Dict
from dataclasses import dataclass, field
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 CleverMicroMessage:
DEFAULT_CONTENT_TYPE = ["application/octet-stream"]
def __init__(
self,
magic: str,
id: SnowflakeId,
m_field: str,
d_field: str,
p_field: str,
headers: Dict[str, List[str]],
trace_info: Dict[str, str],
base64body: bytes,
):
self._magic = magic
self._id: SnowflakeId = id
self._m_field = m_field
self._d_field = d_field
self._p_field = p_field
self._headers: Dict[str, List[str]] = headers
self._trace_info: Dict[str, str] = trace_info
self._base64body: bytes = base64body
def magic(self):
return self._magic
def id(self):
return self._id
def m_field(self) -> str:
return self._m_field
def d_field(self) -> str:
return self._d_field
def p_field(self) -> str:
return self._p_field
def headers(self) -> Dict[str, List[str]]:
return self._headers
def trace_info(self) -> Dict[str, str]:
return self._trace_info
def base64body(self) -> bytes:
return self._base64body
def body(self) -> bytes:
return base64.b64decode(self.base64body())
def contentType(self) -> str:
return self.headers().get("Content-Type", self.DEFAULT_CONTENT_TYPE)[0]
class DataMessage(CleverMicroMessage):
def __init__(
self,
magic: str,
id: SnowflakeId,
method: str,
domain: str,
path: str,
headers: Dict[str, List[str]],
trace_info: Dict[str, str],
base64body: bytes,
):
super().__init__(
magic, id, method, domain, path, headers, trace_info, base64body
)
def method(self) -> str:
return self.m_field()
def domain(self) -> str:
return self.d_field()
def path(self) -> str:
return self.p_field()
class DataResponse(CleverMicroMessage):
def __init__(
self,
magic: str,
id: SnowflakeId,
response_code: str,
error: str,
error_cause: str,
headers: Dict[str, List[str]],
trace_info: Dict[str, str],
base64body: bytes,
):
super().__init__(
magic,
id,
response_code,
error,
error_cause,
headers,
trace_info,
base64body,
)
def response_code(self) -> str:
return self.m_field()
def error(self) -> str:
return self.d_field()
def error_cause(self) -> str:
return self.p_field()
@dataclass(frozen=True)
class AMQRoute:
"""
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__()
@dataclass
class DataMessage:
"""
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.
Please see DataMessageFactory for code to instantiate, serialize and deserialize the class
"""
magic: str
id: SnowflakeId
method: str
domain: str
path: str
headers: dict[str, List[str]]
trace_info: dict[str, str]
base64_body: bytes
def routing_key(self) -> str:
return sanitize_as_rabbitmq_name(f"{self.domain}.{self.path}")
def body(self) -> bytes:
return base64.b64decode(self.base64_body)
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.base64_body,
)
self.extra_data = extra_data
@dataclass
class DataResponse:
"""
The response message to RPC style call (note: the request is made with DataMessage)
Please see DataResponseFactory for code to instantiate, serialize and deserialize the class
"""
id: str
response_code: int
content_type: str
response: bytes
error: str | None
error_cause: str | None
trace_info: dict[str, str]
def body(self) -> bytes:
return base64.b64decode(self.response)
# 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"
SCSLE_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")