Coverage for /Users/eugene/Development/robotnikmq/robotnikmq/core.py: 86%

1 import logging

2 from contextlib import contextmanager

3 from dataclasses import dataclass

4 from datetime import datetime

5 from json import loads as _from_json

6 from json.decoder import JSONDecodeError

7 from pathlib import Path

8 from random import sample

9 from threading import current_thread

10 from typing import Optional, Callable, Any, Dict, Union, Generator, List, TypedDict

11 from uuid import uuid4 as uuid, UUID

12

13 from arrow import Arrow, get as to_arrow, now

14 from funcy import first

15 from pika import BlockingConnection

16 from pika.adapters.blocking_connection import BlockingChannel

17 from pika.exceptions import AMQPError, AMQPConnectionError

18 from pydantic import BaseModel # pylint: disable=E0611

19 from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_random, Retrying

20 from tenacity.after import after_log

21 from typeguard import typechecked

22

23 from robotnikmq.config import RobotnikConfig, config_of, ServerConfig

24 from robotnikmq.error import UnableToConnect, MalformedMessage

25 from robotnikmq.log import log

26 from robotnikmq.utils import to_json as _to_json

27

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

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

30

31

32@contextmanager

33 def 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

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

43 try:

44 _to_json(content)

45 return True

46 except (TypeError, OverflowError):

47 return False

48

49

50@typechecked

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

52 try:

53 _from_json(string)

54 return True

55 except JSONDecodeError:

56 return False

57

58@typechecked

59 class 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)

67 class 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

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