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.
142 lines
4.4 KiB
Python
142 lines
4.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
|
|
mode: str | None = None
|
|
partitions: Any = "all"
|
|
offset_store: dict[str, Any] = field(default_factory=dict)
|
|
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,
|
|
mode: str | None = None,
|
|
partitions: Any = "all",
|
|
offset_store: dict[str, Any] | 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")
|
|
mode = _normalize_mode(mode) if mode is not None else None
|
|
if mode == "assign" and partitions != "all" and not partitions:
|
|
raise ValueError("mq assign consumer partitions 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,
|
|
mode=mode,
|
|
partitions=partitions,
|
|
offset_store=dict(offset_store or {}),
|
|
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)
|
|
|
|
|
|
def _normalize_mode(mode: str) -> str:
|
|
mode = str(mode or "subscribe").strip().lower()
|
|
if mode not in {"subscribe", "assign"}:
|
|
raise ValueError("mq consumer mode must be 'subscribe' or 'assign'")
|
|
return mode
|
|
|
|
|
|
mq_registry = MQRegistry()
|
|
|
|
|
|
def mq_consumer(
|
|
*topics: str,
|
|
name: str | None = None,
|
|
group_id: str | None = None,
|
|
mode: str | None = None,
|
|
partitions: Any = "all",
|
|
offset_store: dict[str, Any] | 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,
|
|
mode=mode,
|
|
partitions=partitions,
|
|
offset_store=offset_store,
|
|
handler=func,
|
|
value_format=value_format,
|
|
failure_backoff_seconds=failure_backoff_seconds,
|
|
config=config,
|
|
)
|
|
return func
|
|
|
|
return decorator
|