amq-adapter-python/amqp/model/model.py

import base64

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 AMQMessage:
    """
    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 AMQResponse class.
    Please see AMQMessageFactory 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 MultipartAMQMessage(AMQMessage):
    def __init__(self, message: AMQMessage, extra_data: dict):
        super().__init__(message.magic, message.id,
                         message.method, message.domain, message.path,
                         message.headers, message.trace_info,
                         message.base64_body)
        self.extra_data = extra_data

@dataclass
class AMQResponse:
    """
    The response message to RPC style call (note: the request is made with AMQMessage)
    Please see AMQResponseFactory 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)


@dataclass(frozen=True)
class CleverMicroServiceMessage:
    """
    The CleverMicro service message that is used to communicate the internal state of the AMQ channel.
    In v0.1 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