Coverage for /Users/eugene/Development/robotnikmq/robotnikmq/core.py: 86%
146 statements
« prev ^ index » next coverage.py v7.3.4, created at 2023-12-26 23:36 -0500
« prev ^ index » next coverage.py v7.3.4, created at 2023-12-26 23:36 -0500
1import logging
2from contextlib import contextmanager
3from dataclasses import dataclass
4from datetime import datetime
5from json import loads as _from_json
6from json.decoder import JSONDecodeError
7from pathlib import Path
8from random import sample
9from threading import current_thread
10from typing import Optional, Callable, Any, Dict, Union, Generator, List, TypedDict
11from uuid import uuid4 as uuid, UUID
13from arrow import Arrow, get as to_arrow, now
14from funcy import first
15from pika import BlockingConnection
16from pika.adapters.blocking_connection import BlockingChannel
17from pika.exceptions import AMQPError, AMQPConnectionError
18from pydantic import BaseModel # pylint: disable=E0611
19from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_random, Retrying
20from tenacity.after import after_log
21from typeguard import typechecked
23from robotnikmq.config import RobotnikConfig, config_of, ServerConfig
24from robotnikmq.error import UnableToConnect, MalformedMessage
25from robotnikmq.log import log
26from robotnikmq.utils import to_json as _to_json
28AMQPErrorCallback = Optional[Callable[[AMQPError], None]]
29ConnErrorCallback = Optional[Callable[[AMQPConnectionError], None]]
32@contextmanager
33def thread_name(name: Union[str, UUID]):
34 thread = current_thread()
35 original = thread.name
36 thread.name = str(name)
37 yield
38 thread.name = original
41@typechecked
42def jsonable(content: Any) -> bool:
43 try:
44 _to_json(content)
45 return True
46 except (TypeError, OverflowError):
47 return False
50@typechecked
51def valid_json(string: str) -> bool:
52 try:
53 _from_json(string)
54 return True
55 except JSONDecodeError:
56 return False
58@typechecked
59class MessageTypedDict(TypedDict):
60 contents: Dict[str, Any]
61 routing_key: Optional[str]
62 timestamp: Union[int, float]
63 msg_id: str
65@typechecked
66@dataclass(frozen=True)
67class Message:
68 contents: Union[BaseModel, Dict[str, Any]]
69 routing_key: str
70 timestamp: Arrow
71 msg_id: Union[str, UUID]
73 @staticmethod
74 def of(
75 contents: Union[BaseModel, Dict],
76 routing_key: Optional[str] = None,
77 timestamp: Union[int, float, datetime, Arrow, None] = None,
78 msg_id: Union[str, UUID, None] = None,
79 ) -> 'Message':
80 msg_id = msg_id or uuid()
81 if not jsonable(contents): 81 ↛ 82, 81 ↛ 832 missed branches: 1) line 81 didn't jump to line 82, because the condition on line 81 was never true, 2) line 81 didn't jump to line 83, because the condition on line 81 was never false
82 raise ValueError("Contents of message have to be JSON-serializeable")
83 contents = contents.dict() if isinstance(contents, BaseModel) else contents
84 routing_key: str = routing_key or ""
85 timestamp: Arrow = to_arrow(timestamp) if timestamp is not None else now()
86 return Message(contents, routing_key, timestamp, msg_id)
88 def with_routing_key(self, routing_key: Optional[str]) -> 'Message':
89 return Message.of(self.contents, routing_key, self.timestamp, self.msg_id)
91 def to_dict(self) -> MessageTypedDict:
92 return {
93 "routing_key": self.routing_key,
94 "contents": self.contents,
95 "msg_id": str(self.msg_id),
96 "timestamp": self.timestamp.int_timestamp,
97 }
99 def to_json(self) -> str:
100 return _to_json(self.to_dict())
102 @staticmethod
103 def of_json(body: str) -> "Message": # pylint: disable=C0103
104 try:
105 msg = _from_json(body)
106 return Message.of(
107 msg["contents"], msg["routing_key"], msg["timestamp"], msg["msg_id"]
108 )
109 except (JSONDecodeError, KeyError) as exc:
110 raise MalformedMessage(body) from exc
112 def __getitem__(self, key: str) -> Any:
113 return self.contents[key]
115 def keys(self):
116 return self.contents.keys()
118 def values(self):
119 return self.contents.values()
121 def __contains__(self, item: str) -> bool:
122 return item in self.contents
124 def __iter__(self):
125 return iter(self.contents)
127 @property
128 def route(self) -> str:
129 return self.routing_key
132class Robotnik:
133 @typechecked
134 def __init__(
135 self,
136 config: Optional[RobotnikConfig] = None,
137 config_paths: Optional[List[Path]] = None,
138 ):
139 config_paths = config_paths or [
140 Path.cwd() / "robotnikmq.yaml",
141 Path.home() / ".config" / "robotnikmq" / "robotnikmq.yaml",
142 Path("/etc/robotnikmq/robotnikmq.yaml"),
143 ]
144 self.config = config or config_of( 144 ↛ exitline 144 didn't jump to the function exit
145 first(path for path in config_paths if path.exists())
146 )
147 self._connection = None
148 self._channel: Optional[BlockingChannel] = None
149 self.log = log.bind(rmq_server="")
151 @retry(stop=stop_after_attempt(3), wait=wait_random(0, 2),
152 reraise=True, after=after_log(log, logging.WARN))
153 @typechecked
154 def _connect_to_server(self, config: ServerConfig) -> BlockingConnection:
155 connection = BlockingConnection(config.conn_params())
156 self.log = log.bind(rmq_server=f"{config.host}:{config.port}{config.vhost}")
157 self.log.success(f"Connection to {config.host}:{config.port}{config.vhost} is successful")
158 return connection
160 @typechecked
161 def _connect_to_cluster(self) -> BlockingConnection:
162 self.log = log.bind(rmq_server="")
163 for tier in self.config.tiers: 163 ↛ 169line 163 didn't jump to line 169, because the loop on line 163 didn't complete
164 for config in sample(tier, len(tier)):
165 try:
166 return self._connect_to_server(config)
167 except AMQPError:
168 log.warning(f"Unable to connect to {config.host}:{config.port}{config.vhost}")
169 raise UnableToConnect("Cannot connect to any of the configured servers")
171 @property
172 def connection(self) -> BlockingConnection:
173 if self._connection is None or not self._connection.is_open:
174 self._connection = self._connect_to_cluster()
175 return self._connection
177 @typechecked
178 def _open_channel(self) -> BlockingChannel:
179 _channel = self.connection.channel()
180 _channel.basic_qos(prefetch_count=1)
181 return _channel
183 @property
184 def channel(self) -> BlockingChannel:
185 if self._channel is None or not self._channel.is_open:
186 self._channel = self._open_channel()
187 return self._channel
189 @contextmanager
190 def open_channel(self) -> Generator[BlockingChannel, None, None]:
191 _ch = self.channel
192 yield _ch
193 self.close_channel(_ch)
195 @typechecked
196 def close_channel(self, channel: Optional[BlockingChannel] = None) -> None:
197 channel = channel or self.channel
198 if channel is not None and channel.is_open: 198 ↛ exitline 198 didn't return from function 'close_channel', because the condition on line 198 was never false
199 channel.stop_consuming()
200 channel.close()