Coverage for documents/plugins/helpers.py: 69%

1import enum 

2from typing import TYPE_CHECKING 

3from typing import Literal 

4from typing import Self 

5from typing import TypeAlias 

6from typing import TypedDict 

7 

8from asgiref.sync import async_to_sync 

9from channels.layers import get_channel_layer 

10 

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

12 from channels_redis.pubsub import RedisPubSubChannelLayer 

13 

14 

15class ProgressStatusOptions(enum.StrEnum): 

16 STARTED = "STARTED" 

17 WORKING = "WORKING" 

18 SUCCESS = "SUCCESS" 

19 FAILED = "FAILED" 

20 

21 

22class PermissionsData(TypedDict, total=False): 

23 """Permission fields included in status messages for access control.""" 

24 

25 owner_id: int | None 

26 users_can_view: list[int] 

27 groups_can_view: list[int] 

28 

29 

30class ProgressUpdateData(TypedDict): 

31 filename: str | None 

32 task_id: str | None 

33 current_progress: int 

34 max_progress: int 

35 status: str 

36 message: str 

37 document_id: int | None 

38 owner_id: int | None 

39 users_can_view: list[int] 

40 groups_can_view: list[int] 

41 

42 

43class StatusUpdatePayload(TypedDict): 

44 type: Literal["status_update"] 

45 data: ProgressUpdateData 

46 

47 

48class DocumentsDeletedData(TypedDict): 

49 documents: list[int] 

50 

51 

52class DocumentsDeletedPayload(TypedDict): 

53 type: Literal["documents_deleted"] 

54 data: DocumentsDeletedData 

55 

56 

57class DocumentUpdatedData(TypedDict): 

58 document_id: int 

59 modified: str 

60 owner_id: int | None 

61 users_can_view: list[int] 

62 groups_can_view: list[int] 

63 

64 

65class DocumentUpdatedPayload(TypedDict): 

66 type: Literal["document_updated"] 

67 data: DocumentUpdatedData 

68 

69 

70WebsocketPayload: TypeAlias = ( 

71 StatusUpdatePayload | DocumentsDeletedPayload | DocumentUpdatedPayload 

72) 

73 

74 

75class BaseStatusManager: 

76 """ 

77 Handles sending of progress information via the channel layer, with proper management 

78 of the open/close of the layer to ensure messages go out and everything is cleaned up 

79 """ 

80 

81 def __init__(self) -> None: 

82 self._channel: RedisPubSubChannelLayer | None = None 

83 

84 def __enter__(self) -> Self: 

85 self.open() 

86 return self 

87 

88 def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: 

89 self.close() 

90 

91 def open(self) -> None: 

92 """ 

93 If not already opened, gets the default channel layer 

94 opened and ready to send messages 

95 """ 

96 if self._channel is None: 96 ↛ exitline 96 didn't return from function 'open' because the condition on line 96 was always true

97 self._channel = get_channel_layer() 

98 

99 def close(self) -> None: 

100 """ 

101 If it was opened, flushes the channel layer 

102 """ 

103 if self._channel is not None: 

104 async_to_sync(self._channel.flush) 

105 self._channel = None 

106 

107 def send(self, payload: WebsocketPayload) -> None: 

108 # Ensure the layer is open 

109 self.open() 

110 

111 # Just for IDEs 

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

113 assert self._channel is not None 

114 

115 # Construct and send the update 

116 async_to_sync(self._channel.group_send)("status_updates", payload) 

117 

118 

119class ProgressManager(BaseStatusManager): 

120 def __init__(self, filename: str | None = None, task_id: str | None = None) -> None: 

121 super().__init__() 

122 self.filename = filename 

123 self.task_id = task_id 

124 

125 def send_progress( 

126 self, 

127 status: ProgressStatusOptions, 

128 message: str, 

129 current_progress: int, 

130 max_progress: int, 

131 *, 

132 document_id: int | None = None, 

133 owner_id: int | None = None, 

134 users_can_view: list[int] | None = None, 

135 groups_can_view: list[int] | None = None, 

136 ) -> None: 

137 data: ProgressUpdateData = { 

138 "filename": self.filename, 

139 "task_id": self.task_id, 

140 "current_progress": current_progress, 

141 "max_progress": max_progress, 

142 "status": status, 

143 "message": message, 

144 "document_id": document_id, 

145 "owner_id": owner_id, 

146 "users_can_view": users_can_view or [], 

147 "groups_can_view": groups_can_view or [], 

148 } 

149 payload: StatusUpdatePayload = {"type": "status_update", "data": data} 

150 self.send(payload) 

151 

152 

153class DocumentsStatusManager(BaseStatusManager): 

154 def send_documents_deleted(self, documents: list[int]) -> None: 

155 payload: DocumentsDeletedPayload = { 

156 "type": "documents_deleted", 

157 "data": { 

158 "documents": documents, 

159 }, 

160 } 

161 self.send(payload) 

162 

163 def send_document_updated( 

164 self, 

165 *, 

166 document_id: int, 

167 modified: str, 

168 owner_id: int | None = None, 

169 users_can_view: list[int] | None = None, 

170 groups_can_view: list[int] | None = None, 

171 ) -> None: 

172 payload: DocumentUpdatedPayload = { 

173 "type": "document_updated", 

174 "data": { 

175 "document_id": document_id, 

176 "modified": modified, 

177 "owner_id": owner_id, 

178 "users_can_view": users_can_view or [], 

179 "groups_can_view": groups_can_view or [], 

180 }, 

181 } 

182 self.send(payload)