Stanislav Hejny
228 lines
8.3 KiB
Python
228 lines
8.3 KiB
Python
import base64
|
|
import json
|
|
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
|
|
from amqp.adapter.file_downloader import download_buffer
|
|
|
|
|
|
class AMQMessageFactory:
|
|
SNOWFLAKE_ID_BITS = 80
|
|
SEQUENCE_BITS = 12
|
|
MACHINE_ID_BITS = 24
|
|
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):
|
|
"""
|
|
return reusable factory instance
|
|
:param machine_id: a generator ID, should be unique for each component to prevent duplicate IDs
|
|
:return: factory instance
|
|
"""
|
|
if cls._instance is None:
|
|
cls._instance = AMQMessageFactory(machine_id)
|
|
return cls._instance
|
|
|
|
def generate_snowflake_id(self) -> SnowflakeId:
|
|
"""
|
|
Generate an unique ID following the SnowflakeID structure, see https://en.wikipedia.org/wiki/Snowflake_ID
|
|
:return: 80-bit wide ID
|
|
"""
|
|
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:
|
|
"""
|
|
Encapsulate / hide the system level inputs, create the AMQ message from business relevant input
|
|
:param method: a method requested (e.g. HTTP method like GET, POST in case of REST request)
|
|
:param domain: service host name
|
|
:param path: path at the service
|
|
:param headers: HTTP headers
|
|
:param message: payload
|
|
:return: AMQMessage object
|
|
"""
|
|
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:
|
|
"""
|
|
Constructs the message specific routing key.
|
|
:param message: AMQMessage input
|
|
:return: routing key (as string) compliant to RabbitMQ allowed character set.
|
|
"""
|
|
return sanitize_as_rabbitmq_name(f"{message.domain}.{message.path}")
|
|
|
|
@staticmethod
|
|
def get_timestamp_from_id(_id: SnowflakeId) -> int:
|
|
"""
|
|
Helper method - retrieve the timestamp from the ID
|
|
:param _id: SnowflakeID (includes encoded timestamp)
|
|
:return: the message timestamp
|
|
"""
|
|
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:
|
|
"""
|
|
Human-readable representation.
|
|
:param message: AMQMessage input
|
|
:return: as string
|
|
"""
|
|
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:
|
|
"""
|
|
Converts AMQMessage to byte array
|
|
:param message: message to serialize
|
|
:return: messages as 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:
|
|
"""
|
|
Builds the AMQMessage from its serialized (byte array) form.
|
|
:param input_bytes: serialized message
|
|
:return: AMQMessage instance
|
|
"""
|
|
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,
|
|
)
|
|
|
|
@staticmethod
|
|
async def from_stream(stream: bytes, rabbitmq_url: str, loop) -> AMQMessage | None:
|
|
"""
|
|
Builds the AMQMessage from its serialized (byte array) form.
|
|
:param stream: serialized message
|
|
:param rabbitmq_url: RabbitMQ connection URL
|
|
:return: AMQMessage 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)
|
|
if eof == 1:
|
|
return AMQMessageFactory.from_bytes(body)
|
|
return None
|