Stanislav Hejny
280 lines
9.8 KiB
Python
280 lines
9.8 KiB
Python
"""
|
|
methods to create, serialize and deserialize DataMessage objects.
|
|
"""
|
|
|
|
import base64
|
|
import json
|
|
import re
|
|
from datetime import datetime, timezone
|
|
from io import BytesIO
|
|
from typing import Any, Dict, List
|
|
|
|
import amqp.adapter.file_handler
|
|
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
|
|
from amqp.adapter.pydantic_serializer import python_type_to_json
|
|
from amqp.adapter.serializer import (
|
|
map_as_string,
|
|
map_with_list_as_string,
|
|
parse_map_list_string,
|
|
parse_map_string,
|
|
)
|
|
from amqp.model.model import DataMessage, DataMessageMagicByte
|
|
from amqp.model.snowflake_id import SnowflakeId
|
|
from amqp.router.utils import sanitize_as_rabbitmq_name
|
|
|
|
|
|
class DataMessageFactory:
|
|
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
|
|
|
|
_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 = DataMessageFactory(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(DataMessageFactory.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,
|
|
) -> DataMessage:
|
|
"""
|
|
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: DataMessage object
|
|
"""
|
|
serializable_headers = headers.copy() if headers else {}
|
|
payload = base64.b64encode(message.encode("utf-8") if isinstance(message, str) else message)
|
|
trace_info = {}
|
|
return DataMessage(
|
|
DataMessageMagicByte.HTTP_REQUEST.value,
|
|
self.generate_snowflake_id(),
|
|
method,
|
|
domain,
|
|
path,
|
|
serializable_headers,
|
|
trace_info,
|
|
payload,
|
|
)
|
|
|
|
def create_rpc_message(
|
|
self,
|
|
fq_method_name: str,
|
|
swarm_id: str,
|
|
args: Dict[str, Any] = None,
|
|
) -> DataMessage:
|
|
"""
|
|
Create a RPC message with the given parameters.
|
|
:param fq_method_name: fully qualified method name (e.g. 'com.example.ServiceClass.method_name')
|
|
:param args: arguments whose serialized form will form the request body
|
|
:return: DataMessage object
|
|
"""
|
|
package, class_name, method = fq_method_name.rsplit(".", 2)
|
|
headers = {"Content-Type": ["application/octet-stream"], "Return-Address": [swarm_id]}
|
|
trace_info = {}
|
|
return DataMessage(
|
|
DataMessageMagicByte.RPC_REQUEST.value,
|
|
self.generate_snowflake_id(),
|
|
package,
|
|
class_name,
|
|
method,
|
|
headers,
|
|
trace_info,
|
|
python_type_to_json(args) if args else b"",
|
|
)
|
|
|
|
@staticmethod
|
|
def amq_message_routing_key(message: DataMessage) -> str:
|
|
"""
|
|
Constructs the message specific routing key.
|
|
:param message: DataMessage 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] >> (DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS)
|
|
high = bits[1] << (
|
|
64 - DataMessageFactory.SEQUENCE_BITS - DataMessageFactory.MACHINE_ID_BITS
|
|
)
|
|
return (high | low) + DataMessageFactory.snowflake_epoch * 1000
|
|
|
|
@staticmethod
|
|
def to_string(message: DataMessage) -> str:
|
|
"""
|
|
Human-readable representation.
|
|
:param message: DataMessage input
|
|
:return: as string
|
|
"""
|
|
return (
|
|
f"DataMessage{{"
|
|
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.base64body().decode()}"
|
|
f"}}"
|
|
)
|
|
|
|
@staticmethod
|
|
def serialize(message: DataMessage) -> bytes:
|
|
"""
|
|
Converts DataMessage 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.base64body())
|
|
with BytesIO(bytearray(msg_length)) as baos:
|
|
baos.write(prolog_len)
|
|
baos.write(DataMessageMagicByte.HTTP_REQUEST.value.encode("utf-8"))
|
|
baos.write(id_bytes)
|
|
baos.write(prolog)
|
|
baos.write(message.base64body())
|
|
return baos.getvalue()
|
|
|
|
@staticmethod
|
|
def from_bytes(input_bytes: bytes) -> DataMessage:
|
|
"""
|
|
Builds the DataMessage from its serialized (byte array) form.
|
|
:param input_bytes: serialized message
|
|
:return: DataMessage instance
|
|
"""
|
|
prolog_len = (input_bytes[0] << 8) + input_bytes[1]
|
|
metadata = input_bytes[2 : prolog_len + 2].decode()
|
|
|
|
pattern = re.compile(
|
|
f"{DataMessageMagicByte.HTTP_REQUEST.value}"
|
|
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"DataMessage 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 DataMessage(
|
|
DataMessageMagicByte.HTTP_REQUEST.value,
|
|
SnowflakeId.from_hex(id_str),
|
|
method,
|
|
domain,
|
|
path,
|
|
headers,
|
|
trace_info,
|
|
body_b64,
|
|
)
|
|
|
|
@staticmethod
|
|
async def from_stream(
|
|
stream: bytes,
|
|
rabbitmq_url: str,
|
|
loop,
|
|
file_handler: FileHandler = StreamingFileHandler(),
|
|
) -> DataMessage | None:
|
|
"""
|
|
Builds the DataMessage from its serialized (byte array) form.
|
|
:param stream: serialized message
|
|
:param rabbitmq_url: RabbitMQ connection URL
|
|
:param loop: event loop
|
|
:param file_handler: file handler for downloading the message
|
|
:return: DataMessage instance
|
|
"""
|
|
# async implementation
|
|
_message_data = json.loads(stream.decode())
|
|
_req_info = _message_data.get("req-info", "")
|
|
if _req_info:
|
|
(body, eof) = await file_handler.download_buffer(
|
|
loop=loop, rabbitmq_url=rabbitmq_url, queue_name=_req_info
|
|
)
|
|
if eof == amqp.adapter.file_handler.END_OF_FILE:
|
|
return DataMessageFactory.from_bytes(body)
|
|
return None
|