Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/utilities/messaging/__init__.py: 78%

1import abc 

2from contextlib import asynccontextmanager, AbstractAsyncContextManager 

3from dataclasses import dataclass 

4import importlib 

5from types import TracebackType 

6from typing import ( 

7 TYPE_CHECKING, 

8 Any, 

9 Callable, 

10 Optional, 

11 Protocol, 

12 Type, 

13 TypeVar, 

14 Union, 

15 runtime_checkable, 

16) 

17from collections.abc import AsyncGenerator, Awaitable, Iterable, Mapping 

18from typing_extensions import Self 

19 

20from prefect.settings import PREFECT_MESSAGING_CACHE, PREFECT_MESSAGING_BROKER 

21from prefect.logging import get_logger 

22 

23if TYPE_CHECKING: 23 ↛ 24line 23 didn't jump to line 24 because the condition on line 23 was never true

24 import logging 

25 

26logger: "logging.Logger" = get_logger(__name__) 

27 

28 

29M = TypeVar("M", bound="Message", covariant=True) 

30 

31 

32class Message(Protocol): 

33 """ 

34 A protocol representing a message sent to a message broker. 

35 """ 

36 

37 @property 

38 def data(self) -> Union[str, bytes]: ... 38 ↛ exitline 38 didn't return from function 'data' because

39 

40 @property 

41 def attributes(self) -> Mapping[str, Any]: ... 41 ↛ exitline 41 didn't return from function 'attributes' because

42 

43 

44class Cache(abc.ABC): 

45 @abc.abstractmethod 

46 async def clear_recently_seen_messages(self) -> None: ... 46 ↛ exitline 46 didn't return from function 'clear_recently_seen_messages' because

47 

48 @abc.abstractmethod 

49 async def without_duplicates( 49 ↛ exitline 49 didn't return from function 'without_duplicates' because

50 self, attribute: str, messages: Iterable[M] 

51 ) -> list[M]: ... 

52 

53 @abc.abstractmethod 

54 async def forget_duplicates( 54 ↛ exitline 54 didn't return from function 'forget_duplicates' because

55 self, attribute: str, messages: Iterable[Message] 

56 ) -> None: ... 

57 

58 

59class Publisher(AbstractAsyncContextManager["Publisher"], abc.ABC): 

60 def __init__( 60 ↛ exitline 60 didn't return from function '__init__' because

61 self, 

62 topic: str, 

63 cache: Optional[Cache] = None, 

64 deduplicate_by: Optional[str] = None, 

65 ) -> None: ... 

66 

67 @abc.abstractmethod 

68 async def publish_data( 68 ↛ exitline 68 didn't return from function 'publish_data' because

69 self, data: bytes, attributes: Mapping[str, str] 

70 ) -> None: ... 

71 

72 @abc.abstractmethod 

73 async def __aenter__(self) -> Self: ... 73 ↛ exitline 73 didn't return from function '__aenter__' because

74 

75 @abc.abstractmethod 

76 async def __aexit__( 76 ↛ exitline 76 didn't return from function '__aexit__' because

77 self, 

78 exc_type: Optional[Type[BaseException]], 

79 exc_val: Optional[BaseException], 

80 exc_tb: Optional[TracebackType], 

81 ) -> None: ... 

82 

83 

84@dataclass 

85class CapturedMessage: 

86 data: bytes 

87 attributes: Mapping[str, str] 

88 

89 

90class CapturingPublisher(Publisher): 

91 messages: list[CapturedMessage] = [] 

92 deduplicate_by: Optional[str] 

93 

94 def __init__( 

95 self, 

96 topic: str, 

97 cache: Optional[Cache] = None, 

98 deduplicate_by: Optional[str] = None, 

99 ) -> None: 

100 self.topic = topic 

101 self.cache: Cache = cache or create_cache() 

102 self.deduplicate_by = deduplicate_by 

103 

104 async def __aenter__(self) -> Self: 

105 return self 

106 

107 async def __aexit__( 

108 self, 

109 exc_type: Optional[Type[BaseException]], 

110 exc_val: Optional[BaseException], 

111 exc_tb: Optional[TracebackType], 

112 ) -> None: 

113 pass 

114 

115 async def publish_data(self, data: bytes, attributes: Mapping[str, str]) -> None: 

116 to_publish = [CapturedMessage(data, attributes)] 

117 

118 if self.deduplicate_by: 

119 to_publish = await self.cache.without_duplicates( 

120 self.deduplicate_by, to_publish 

121 ) 

122 

123 self.messages.extend(to_publish) 

124 

125 

126MessageHandler = Callable[[Message], Awaitable[None]] 

127 

128 

129class StopConsumer(Exception): 

130 """ 

131 Exception to raise to stop a consumer. 

132 """ 

133 

134 def __init__(self, ack: bool = False): 

135 self.ack = ack 

136 

137 

138class Consumer(abc.ABC): 

139 """ 

140 Abstract base class for consumers that receive messages from a message broker and 

141 call a handler function for each message received. 

142 """ 

143 

144 def __init__(self, topic: str, **kwargs: Any) -> None: 

145 self.topic = topic 

146 

147 @abc.abstractmethod 

148 async def run(self, handler: MessageHandler) -> None: 

149 """Runs the consumer (indefinitely)""" 

150 ... 

151 

152 async def cleanup(self) -> None: 

153 """Cleanup resources when the consumer is stopped. 

154 

155 Override this method in subclasses that need to perform cleanup, 

156 such as unsubscribing from topics or closing connections. 

157 

158 The default implementation is a no-op, which is appropriate for 

159 consumers that don't need explicit cleanup. 

160 """ 

161 pass 

162 

163 

164@runtime_checkable 

165class CacheModule(Protocol): 

166 Cache: type[Cache] 

167 

168 

169def create_cache() -> Cache: 

170 """ 

171 Creates a new cache with the applications default settings. 

172 

173 Returns: 

174 a new Cache instance 

175 """ 

176 module = importlib.import_module(PREFECT_MESSAGING_CACHE.value()) 

177 assert isinstance(module, CacheModule) 

178 

179 return module.Cache() 

180 

181 

182@runtime_checkable 

183class BrokerModule(Protocol): 

184 Publisher: type[Publisher] 

185 Consumer: type[Consumer] 

186 ephemeral_subscription: Callable[ 

187 [str], AbstractAsyncContextManager[Mapping[str, Any]] 

188 ] 

189 

190 # Used for testing: a context manager that breaks the topic in a way that raises 

191 # a ValueError("oops") when attempting to publish a message. 

192 break_topic: Callable[[], AbstractAsyncContextManager[None]] 

193 

194 

195def create_publisher( 

196 topic: str, cache: Optional[Cache] = None, deduplicate_by: Optional[str] = None 

197) -> Publisher: 

198 """ 

199 Creates a new publisher with the applications default settings. 

200 Args: 

201 topic: the topic to publish to 

202 Returns: 

203 a new Consumer instance 

204 """ 

205 cache = cache or create_cache() 

206 

207 module = importlib.import_module(PREFECT_MESSAGING_BROKER.value()) 

208 assert isinstance(module, BrokerModule) 

209 return module.Publisher(topic, cache, deduplicate_by=deduplicate_by) 

210 

211 

212@asynccontextmanager 

213async def ephemeral_subscription(topic: str) -> AsyncGenerator[Mapping[str, Any], Any]: 

214 """ 

215 Creates an ephemeral subscription to the given source, removing it when the context 

216 exits. 

217 """ 

218 module = importlib.import_module(PREFECT_MESSAGING_BROKER.value()) 

219 assert isinstance(module, BrokerModule) 

220 async with module.ephemeral_subscription(topic) as consumer_create_kwargs: 

221 yield consumer_create_kwargs 

222 

223 

224def create_consumer(topic: str, **kwargs: Any) -> Consumer: 

225 """ 

226 Creates a new consumer with the applications default settings. 

227 Args: 

228 topic: the topic to consume from 

229 Returns: 

230 a new Consumer instance 

231 """ 

232 module = importlib.import_module(PREFECT_MESSAGING_BROKER.value()) 

233 assert isinstance(module, BrokerModule) 

234 return module.Consumer(topic, **kwargs)