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

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 

12 

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 

22 

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 

27 

28AMQPErrorCallback = Optional[Callable[[AMQPError], None]] 

29ConnErrorCallback = Optional[Callable[[AMQPConnectionError], None]] 

30 

31 

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 

39 

40 

41@typechecked 

42def jsonable(content: Any) -> bool: 

43 try: 

44 _to_json(content) 

45 return True 

46 except (TypeError, OverflowError): 

47 return False 

48 

49 

50@typechecked 

51def valid_json(string: str) -> bool: 

52 try: 

53 _from_json(string) 

54 return True 

55 except JSONDecodeError: 

56 return False 

57 

58@typechecked 

59class MessageTypedDict(TypedDict): 

60 contents: Dict[str, Any] 

61 routing_key: Optional[str] 

62 timestamp: Union[int, float] 

63 msg_id: str 

64 

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] 

72 

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) 

87 

88 def with_routing_key(self, routing_key: Optional[str]) -> 'Message': 

89 return Message.of(self.contents, routing_key, self.timestamp, self.msg_id) 

90 

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 } 

98 

99 def to_json(self) -> str: 

100 return _to_json(self.to_dict()) 

101 

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 

111 

112 def __getitem__(self, key: str) -> Any: 

113 return self.contents[key] 

114 

115 def keys(self): 

116 return self.contents.keys() 

117 

118 def values(self): 

119 return self.contents.values() 

120 

121 def __contains__(self, item: str) -> bool: 

122 return item in self.contents 

123 

124 def __iter__(self): 

125 return iter(self.contents) 

126 

127 @property 

128 def route(self) -> str: 

129 return self.routing_key 

130 

131 

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="") 

150 

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 

159 

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") 

170 

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 

176 

177 @typechecked 

178 def _open_channel(self) -> BlockingChannel: 

179 _channel = self.connection.channel() 

180 _channel.basic_qos(prefetch_count=1) 

181 return _channel 

182 

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 

188 

189 @contextmanager 

190 def open_channel(self) -> Generator[BlockingChannel, None, None]: 

191 _ch = self.channel 

192 yield _ch 

193 self.close_channel(_ch) 

194 

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()