Coverage for /usr/local/lib/python3.10/site-packages/opal_common-0.0.0-py3.10.egg/opal_common/topics/publisher.py: 57%

1import asyncio 

2from typing import Any, Optional, Set 

3 

4from ddtrace import tracer 

5from fastapi_websocket_pubsub import PubSubClient, PubSubEndpoint, Topic, TopicList 

6from opal_common.logger import logger 

7 

8 

9class TopicPublisher: 

10 """Abstract publisher, base class for client side and server side 

11 publisher.""" 

12 

13 def __init__(self): 

14 """Inits the publisher's asyncio tasks list.""" 

15 self._tasks: Set[asyncio.Task] = set() 

16 self._tasks_lock = asyncio.Lock() 

17 

18 async def publish(self, topics: TopicList, data: Any = None): 

19 raise NotImplementedError() 

20 

21 async def __aenter__(self): 

22 self.start() 

23 return self 

24 

25 async def __aexit__(self, exc_type, exc, tb): 

26 await self.stop() 

27 

28 def start(self): 

29 """Starts the publisher.""" 

30 logger.debug("started topic publisher") 

31 

32 async def _add_task(self, task: asyncio.Task): 

33 async with self._tasks_lock: 

34 self._tasks.add(task) 

35 task.add_done_callback(self._cleanup_task) 

36 

37 async def wait(self): 

38 async with self._tasks_lock: 

39 await asyncio.gather(*self._tasks, return_exceptions=True) 

40 self._tasks.clear() 

41 

42 async def stop(self): 

43 """Stops the publisher (cancels any running publishing tasks)""" 

44 logger.debug("stopping topic publisher") 

45 await self.wait() 

46 

47 def _cleanup_task(self, task: asyncio.Task): 

48 try: 

49 self._tasks.remove(task) 

50 except KeyError: 

51 ... 

52 

53 

54class PeriodicPublisher: 

55 """Wrapper for a task that publishes to topic on fixed interval 

56 periodically.""" 

57 

58 def __init__( 

59 self, 

60 publisher: TopicPublisher, 

61 time_interval: int, 

62 topic: Topic, 

63 message: Any = None, 

64 task_name: str = "periodic publish task", 

65 ): 

66 """Inits the publisher. 

67 

68 Args: 

69 publisher (TopicPublisher): can publish messages on the pub/sub channel 

70 interval (int): the time interval between publishing consecutive messages 

71 topic (Topic): the topic to publish on 

72 message (Any): the message to publish 

73 """ 

74 self._publisher = publisher 

75 self._interval = time_interval 

76 self._topic = topic 

77 self._message = message 

78 self._task_name = task_name 

79 self._task: Optional[asyncio.Task] = None 

80 

81 async def __aenter__(self): 

82 self.start() 

83 return self 

84 

85 async def __aexit__(self, exc_type, exc, tb): 

86 await self.stop() 

87 

88 def start(self): 

89 """Starts the periodic publisher task.""" 

90 if self._task is not None: 

91 logger.warning(f"{self._task_name} already started") 

92 return 

93 

94 logger.info( 

95 f"started {self._task_name}: topic is '{self._topic}', interval is {self._interval} seconds" 

96 ) 

97 self._task = asyncio.create_task(self._publish_task()) 

98 

99 async def stop(self): 

100 """Stops the publisher (cancels any running publishing tasks)""" 

101 if self._task is not None: 

102 self._task.cancel() 

103 try: 

104 await self._task 

105 except asyncio.CancelledError: 

106 pass 

107 self._task = None 

108 logger.info(f"cancelled {self._task_name} to topic: {self._topic}") 

109 

110 async def wait_until_done(self): 

111 await self._task 

112 

113 async def _publish_task(self): 

114 while True: 

115 await asyncio.sleep(self._interval) 

116 logger.info( 

117 f"{self._task_name}: publishing message on topic '{self._topic}', next publish is scheduled in {self._interval} seconds" 

118 ) 

119 async with self._publisher: 

120 await self._publisher.publish(topics=[self._topic], data=self._message) 

121 

122 

123class ServerSideTopicPublisher(TopicPublisher): 

124 """A simple wrapper around a PubSubEndpoint that exposes publish().""" 

125 

126 def __init__(self, endpoint: PubSubEndpoint): 

127 """Inits the publisher. 

128 

129 Args: 

130 endpoint (PubSubEndpoint): a pub/sub endpoint 

131 """ 

132 self._endpoint = endpoint 

133 super().__init__() 

134 

135 async def _publish_impl(self, topics: TopicList, data: Any = None): 

136 with tracer.trace("topic_publisher.publish", resource=str(topics)): 

137 await self._endpoint.publish(topics=topics, data=data) 

138 

139 async def publish(self, topics: TopicList, data: Any = None): 

140 await self._add_task(asyncio.create_task(self._publish_impl(topics, data))) 

141 

142 

143class ClientSideTopicPublisher(TopicPublisher): 

144 """A simple wrapper around a PubSubClient that exposes publish(). 

145 

146 Provides start() and stop() shortcuts that helps treat this client 

147 as a separate "process" or task that runs in the background. 

148 """ 

149 

150 def __init__(self, client: PubSubClient, server_uri: str): 

151 """Inits the publisher. 

152 

153 Args: 

154 client (PubSubClient): a configured not-yet-started pub sub client 

155 server_uri (str): the URI of the pub sub server we publish to 

156 """ 

157 self._client = client 

158 self._server_uri = server_uri 

159 super().__init__() 

160 

161 def start(self): 

162 """Starts the pub/sub client as a background asyncio task. 

163 

164 the client will attempt to connect to the pubsub server until 

165 successful. 

166 """ 

167 super().start() 

168 self._client.start_client(f"{self._server_uri}") 

169 

170 async def stop(self): 

171 """Stops the pubsub client, and cancels any publishing tasks.""" 

172 await self._client.disconnect() 

173 await super().stop() 

174 

175 async def wait_until_done(self): 

176 """When the publisher is a used as a context manager, this method waits 

177 until the client is done (i.e: terminated) to prevent exiting the 

178 context.""" 

179 return await self._client.wait_until_done() 

180 

181 async def publish(self, topics: TopicList, data: Any = None): 

182 """Publish a message by launching a background task on the event loop. 

183 

184 Args: 

185 topics (TopicList): a list of topics to publish the message to 

186 data (Any): optional data to publish as part of the message 

187 """ 

188 await self._add_task( 

189 asyncio.create_task(self._publish(topics=topics, data=data)) 

190 ) 

191 

192 async def _publish(self, topics: TopicList, data: Any = None) -> bool: 

193 """Do not trigger directly, must be triggered via publish() in order to 

194 run as a monitored background asyncio task.""" 

195 await self._client.wait_until_ready() 

196 logger.info("Publishing to topics: {topics}", topics=topics) 

197 return await self._client.publish(topics, data) 

198 

199 

200class ScopedServerSideTopicPublisher(ServerSideTopicPublisher): 

201 def __init__(self, endpoint: PubSubEndpoint, scope_id: str): 

202 super().__init__(endpoint) 

203 self._scope_id = scope_id 

204 

205 async def publish(self, topics: TopicList, data: Any = None): 

206 scoped_topics = [f"{self._scope_id}:{topic}" for topic in topics] 

207 logger.info("Publishing to topics: {topics}", topics=scoped_topics) 

208 await super().publish(scoped_topics, data)