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.
117 lines
3.4 KiB
Python
117 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MQProducerDefinition:
|
|
name: str
|
|
topic: str
|
|
value_format: str = "json"
|
|
config: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MQConsumerDefinition:
|
|
name: str
|
|
topics: tuple[str, ...]
|
|
handler: Callable
|
|
group_id: str | None = None
|
|
value_format: str = "json"
|
|
failure_backoff_seconds: float | None = None
|
|
config: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class MQRegistry:
|
|
producers: dict[str, MQProducerDefinition] = field(default_factory=dict)
|
|
consumers: dict[str, MQConsumerDefinition] = field(default_factory=dict)
|
|
|
|
def register_producer(
|
|
self,
|
|
*,
|
|
name: str,
|
|
topic: str,
|
|
value_format: str = "json",
|
|
config: dict[str, Any] | None = None,
|
|
) -> MQProducerDefinition:
|
|
if not name:
|
|
raise ValueError("mq producer name is required")
|
|
if not topic:
|
|
raise ValueError("mq producer topic is required")
|
|
if name in self.producers:
|
|
raise ValueError(f"mq producer already registered: {name}")
|
|
definition = MQProducerDefinition(
|
|
name=name,
|
|
topic=topic,
|
|
value_format=value_format,
|
|
config=dict(config or {}),
|
|
)
|
|
self.producers[name] = definition
|
|
return definition
|
|
|
|
def register_consumer(
|
|
self,
|
|
*,
|
|
name: str,
|
|
topics: list[str] | tuple[str, ...] | str,
|
|
handler: Callable,
|
|
group_id: str | None = None,
|
|
value_format: str = "json",
|
|
failure_backoff_seconds: float | None = None,
|
|
config: dict[str, Any] | None = None,
|
|
) -> MQConsumerDefinition:
|
|
if not name:
|
|
raise ValueError("mq consumer name is required")
|
|
topic_values = _normalize_topics(topics)
|
|
if not topic_values:
|
|
raise ValueError("mq consumer topics are required")
|
|
if name in self.consumers:
|
|
raise ValueError(f"mq consumer already registered: {name}")
|
|
definition = MQConsumerDefinition(
|
|
name=name,
|
|
topics=topic_values,
|
|
handler=handler,
|
|
group_id=group_id,
|
|
value_format=value_format,
|
|
failure_backoff_seconds=failure_backoff_seconds,
|
|
config=dict(config or {}),
|
|
)
|
|
self.consumers[name] = definition
|
|
return definition
|
|
|
|
|
|
def _normalize_topics(topics: list[str] | tuple[str, ...] | str) -> tuple[str, ...]:
|
|
if isinstance(topics, str):
|
|
topics = (topics,)
|
|
return tuple(topic for topic in topics if topic)
|
|
|
|
|
|
mq_registry = MQRegistry()
|
|
|
|
|
|
def mq_consumer(
|
|
*topics: str,
|
|
name: str | None = None,
|
|
group_id: str | None = None,
|
|
value_format: str = "json",
|
|
failure_backoff_seconds: float | None = None,
|
|
config: dict[str, Any] | None = None,
|
|
):
|
|
def decorator(func: Callable) -> Callable:
|
|
consumer_name = name or ".".join(topics) or func.__name__
|
|
mq_registry.register_consumer(
|
|
name=consumer_name,
|
|
topics=topics,
|
|
group_id=group_id,
|
|
handler=func,
|
|
value_format=value_format,
|
|
failure_backoff_seconds=failure_backoff_seconds,
|
|
config=config,
|
|
)
|
|
return func
|
|
|
|
return decorator
|