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

import base64
import re
from io import BytesIO
from datetime import datetime, timezone
from typing import Dict, List

from amqp.model.model import AMQMessage
from amqp.model.snowflake_id import SnowflakeId
from amqp.router.serializer import map_with_list_as_string, map_as_string, parse_map_list_string, parse_map_string
from amqp.router.utils import sanitize_as_rabbitmq_name


class AMQMessageFactory:
    SNOWFLAKE_ID_BITS = 80
    SEQUENCE_BITS = 12
    MACHINE_ID_BITS = 26
    TIMESTAMP_BITS = SNOWFLAKE_ID_BITS - SEQUENCE_BITS - MACHINE_ID_BITS

    MACHINE_ID_MASK = (1 << MACHINE_ID_BITS) - 1
    SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1
    MAGIC_BYTE_HTTP_REQUEST = "A"

    _instance = None
    snowflake_epoch = datetime(2024, 10, 1, 0, 0, tzinfo=timezone.utc).timestamp()

    def __init__(self, machine_id: int):
        self.machine_id = (machine_id & self.MACHINE_ID_MASK) << self.SEQUENCE_BITS
        self.snowflake_sequence = 0
        self.snowflake_last_timestamp = 0

    @classmethod
    def get_instance(cls, machine_id: int):
        if cls._instance is None:
            cls._instance = AMQMessageFactory(machine_id)
        return cls._instance

    def generate_snowflake_id(self) -> SnowflakeId:
        timestamp = int(datetime.now(timezone.utc).timestamp() * 1000)

        if timestamp < self.snowflake_last_timestamp:
            raise RuntimeError("Clock moved backwards. Refusing to generate id")

        if timestamp == self.snowflake_last_timestamp:
            self.snowflake_sequence = (self.snowflake_sequence + 1) & self.SEQUENCE_MASK
            if self.snowflake_sequence == 0:
                while timestamp == self.snowflake_last_timestamp:
                    timestamp = int(datetime.now(timezone.utc).timestamp() * 1000)
                self.snowflake_last_timestamp = timestamp
        else:
            self.snowflake_sequence = 0
            self.snowflake_last_timestamp = timestamp

        time = (timestamp - 1000 * int(AMQMessageFactory.snowflake_epoch))
        return SnowflakeId(
            time >> (64 - self.SEQUENCE_BITS - self.MACHINE_ID_BITS),
            ((time << (self.SEQUENCE_BITS + self.MACHINE_ID_BITS)) | self.machine_id | self.snowflake_sequence) & 0xFFFFFFFFFFFFFFFF
        )

    def create_request_message(
            self,
            method: str,
            domain: str,
            path: str,
            headers: Dict[str, List[str]],
            message,
    ) -> AMQMessage:
        serializable_headers = headers.copy() if headers else {}
        payload = base64.b64encode(message.encode('utf-8') if isinstance(message, str) else message)
        trace_info = {}
        return AMQMessage(
            self.MAGIC_BYTE_HTTP_REQUEST,
            self.generate_snowflake_id(),
            method,
            domain,
            path,
            serializable_headers,
            trace_info,
            payload,
        )

    @staticmethod
    def amq_message_routing_key(message: AMQMessage) -> str:
        return sanitize_as_rabbitmq_name(f"{message.domain}.{message.path}")

    @staticmethod
    def get_timestamp_from_id(_id: SnowflakeId) -> int:
        bits = _id.get_bits()
        low = bits[0] >> (AMQMessageFactory.SEQUENCE_BITS + AMQMessageFactory.MACHINE_ID_BITS)
        high = bits[1] << (64 - AMQMessageFactory.SEQUENCE_BITS - AMQMessageFactory.MACHINE_ID_BITS)
        return (high | low) + AMQMessageFactory.snowflake_epoch * 1000

    @staticmethod
    def to_string(message: AMQMessage) -> str:
        return (
            f"AMQMessage{{"
            f"id={message.id}, "
            f"method='{message.method}', "
            f"domain='{message.domain}', "
            f"path='{message.path}', "
            f"headers={{{map_with_list_as_string(message.headers)}}}, "
            f"traceInfo={{{map_as_string(message.trace_info)}}}, "
            f"base64body={message.base64_body.decode()}"
            f"}}"
        )

    @staticmethod
    def serialize(message: AMQMessage) -> bytes:
        id_bytes = str(message.id).encode()
        prolog = (
            f"|m={message.method}|d={message.domain}|p={message.path}|h={{{map_with_list_as_string(message.headers)}}}|t={{{map_as_string(message.trace_info)}}}|b="
        ).encode('utf-8')
        header_length = 1 + len(id_bytes) + len(prolog)
        prolog_len = header_length.to_bytes(2, byteorder='big')
        msg_length = 2 + header_length + len(message.base64_body)
        with BytesIO(bytearray(msg_length)) as baos:
            baos.write(prolog_len)
            baos.write(AMQMessageFactory.MAGIC_BYTE_HTTP_REQUEST.encode('utf-8'))
            baos.write(id_bytes)
            baos.write(prolog)
            baos.write(message.base64_body)
            return baos.getvalue()

    @staticmethod
    def from_bytes(input_bytes: bytes) -> AMQMessage:
        prolog_len = (input_bytes[0] << 8) + input_bytes[1]
        metadata = input_bytes[2:prolog_len + 2].decode()

        pattern = re.compile(
            f"{AMQMessageFactory.MAGIC_BYTE_HTTP_REQUEST}"
            r"([0-9A-Fa-f]+)\|"
            r"m=([^|]*)\|"
            r"d=([^|]*)\|"
            r"p=([^|]*)\|"
            r"h=(\{[^}]*\})\|"
            r"t=(\{[^}]*\})\|"
            r"b="
        )
        match = pattern.match(metadata)
        if not match:
            raise ValueError(
                f"AMQMessage format does not match expected format. "
                f"Input string [{metadata}] does not match pattern: {pattern.pattern}"
            )

        id_str = match.group(1)
        method = match.group(2)
        domain = match.group(3)
        path = match.group(4)
        headers_str = match.group(5)
        trace_info_str = match.group(6)
        body_b64 = input_bytes[prolog_len + 2:]

        headers = parse_map_list_string(headers_str)
        trace_info = parse_map_string(trace_info_str)

        return AMQMessage(
            AMQMessageFactory.MAGIC_BYTE_HTTP_REQUEST,
            SnowflakeId.from_hex(id_str),
            method,
            domain,
            path,
            headers,
            trace_info,
            body_b64,
        )