You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
111 lines
3.0 KiB
Python
111 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from .errors import MQConfigError
|
|
from .registry import MQRegistry
|
|
from .serialization import encode_message_key, encode_message_value
|
|
|
|
|
|
class MQClient:
|
|
def __init__(self, producer: Any, registry: MQRegistry) -> None:
|
|
self._producer = producer
|
|
self._registry = registry
|
|
|
|
def send_json(
|
|
self,
|
|
topic: str,
|
|
value: Any,
|
|
*,
|
|
key: str | bytes | None = None,
|
|
headers: dict[str, str | bytes | None] | None = None,
|
|
) -> None:
|
|
self.send(topic, value=value, key=key, headers=headers, value_format="json")
|
|
|
|
def send(
|
|
self,
|
|
topic: str,
|
|
*,
|
|
value: Any,
|
|
key: str | bytes | None = None,
|
|
headers: dict[str, str | bytes | None] | None = None,
|
|
value_format: str = "bytes",
|
|
) -> None:
|
|
self._producer.poll(0)
|
|
self._producer.produce(
|
|
topic,
|
|
value=encode_message_value(value, value_format),
|
|
key=encode_message_key(key),
|
|
headers=_encode_headers(headers),
|
|
)
|
|
|
|
def sender(self, name: str) -> "MQSender":
|
|
definition = self._registry.producers.get(name)
|
|
if definition is None:
|
|
raise MQConfigError(f"mq producer not registered: {name}")
|
|
return MQSender(self, definition.topic, definition.value_format)
|
|
|
|
def flush(self, timeout: float | None = None) -> None:
|
|
if timeout is None:
|
|
self._producer.flush()
|
|
else:
|
|
self._producer.flush(timeout)
|
|
|
|
|
|
class MQSender:
|
|
def __init__(self, client: MQClient, topic: str, value_format: str) -> None:
|
|
self._client = client
|
|
self.topic = topic
|
|
self.value_format = value_format
|
|
|
|
def send_json(
|
|
self,
|
|
value: Any,
|
|
*,
|
|
key: str | bytes | None = None,
|
|
headers: dict[str, str | bytes | None] | None = None,
|
|
) -> None:
|
|
self._client.send(
|
|
self.topic,
|
|
value=value,
|
|
key=key,
|
|
headers=headers,
|
|
value_format="json",
|
|
)
|
|
|
|
def send(
|
|
self,
|
|
value: Any,
|
|
*,
|
|
key: str | bytes | None = None,
|
|
headers: dict[str, str | bytes | None] | None = None,
|
|
) -> None:
|
|
self._client.send(
|
|
self.topic,
|
|
value=value,
|
|
key=key,
|
|
headers=headers,
|
|
value_format=self.value_format,
|
|
)
|
|
|
|
|
|
def _encode_headers(
|
|
headers: dict[str, str | bytes | None] | None,
|
|
) -> list[tuple[str, bytes | None]] | None:
|
|
if headers is None:
|
|
return None
|
|
result: list[tuple[str, bytes | None]] = []
|
|
for key, value in headers.items():
|
|
if value is None or isinstance(value, bytes):
|
|
result.append((key, value))
|
|
else:
|
|
result.append((key, str(value).encode("utf-8")))
|
|
return result
|
|
|
|
|
|
def mq_client(app) -> MQClient:
|
|
client = getattr(app.state, "iti_mq_client", None)
|
|
if client is None:
|
|
raise MQConfigError("mq client is not configured")
|
|
return client
|