Coverage for gws-app/gws/plugin/upload_helper/__init__.py: 95%

1"""Helper for chunked file uploads. 

2 

3Large files are uploaded by the client in chunks, one request per chunk. The 

4helper stores the chunks in the ephemeral directory and, when asked for the 

5upload, joins them into one file. The helper does not provide endpoints of its 

6own: an action declares an endpoint that receives the chunks and passes them to 

7the helper, and another endpoint that processes the finished upload. 

8 

9Each chunk request contains the file name, the total size, the number of chunks 

10and the chunk number, starting from 0. The first chunk has an empty 

11``uploadUid``, which starts a new upload. The handler responds with the 

12``uploadUid``, which subsequent chunks must provide. Chunks can come in any 

13order. The total size and the number of chunks are limited by ``maxSize``: the 

14maximum number of chunks is derived from it, so that a chunk is at least 500 KB 

15on average. 

16 

17Once the client has sent all chunks, it calls another endpoint of the action 

18with the ``uploadUid``. That endpoint calls ``get_upload`` to get the final 

19file. The file is stored in a temporary location and should be moved to a 

20permanent location if necessary. 

21 

22Example:: 

23 

24 helpers+ { type "upload" maxSize 2000 } 

25 

26Example:: 

27 

28 import gws.plugin.upload_helper as uh 

29 

30 

31 @gws.ext.command.api('myUpload') 

32 def do_upload(self, req, p: uh.ChunkRequest) -> uh.ChunkResponse: 

33 # check permissions, etc... 

34 helper = self.root.app.helper('upload') 

35 return helper.handle_chunk_request(req, p) 

36 

37 @gws.ext.command.api('myProcessUploadedFile') 

38 def do_process(self, req, p: MyProcessRequest): 

39 helper = self.root.app.helper('upload') 

40 try: 

41 upload = helper.get_upload(p.uploadUid) 

42 except uh.Error: 

43 ...upload not ready yet... 

44 ...process(upload.path) 

45""" 

46 

47import shutil 

48 

49import gws 

50import gws.lib.jsonx 

51import gws.lib.osx 

52 

53 

54@gws.ext.config.helper('upload') 

55class Config(gws.Config): 

56 """Helper that receives chunked file uploads.""" 

57 

58 maxSize: int = 1000 

59 """Maximum upload size in megabytes.""" 

60 

61 

62class ChunkRequest(gws.Request): 

63 """Request that carries one chunk of an upload.""" 

64 

65 uploadUid: str = '' 

66 """Upload uid returned for the first chunk, empty for the first chunk.""" 

67 fileName: str 

68 """Name of the uploaded file.""" 

69 totalSize: int 

70 """Total size of the file in bytes.""" 

71 chunkNumber: int 

72 """Number of this chunk, starting from 0.""" 

73 chunkCount: int 

74 """Total number of chunks.""" 

75 content: bytes 

76 """Chunk content.""" 

77 

78 

79class ChunkResponse(gws.Response): 

80 """Response to a chunk request.""" 

81 

82 uploadUid: str 

83 """Upload uid, to be passed with subsequent chunks.""" 

84 

85 

86class Upload(gws.Data): 

87 """State of an upload.""" 

88 

89 uid: str 

90 """Upload uid.""" 

91 fileName: str 

92 """Name of the uploaded file, as sent by the client.""" 

93 totalSize: int 

94 """Total size of the file in bytes.""" 

95 chunkCount: int 

96 """Total number of chunks.""" 

97 path: str 

98 """Path of the assembled file, empty until the upload is finalized.""" 

99 

100 

101class Error(gws.Error): 

102 """Upload error, raised for invalid chunks and incomplete or missing uploads.""" 

103 

104 pass 

105 

106 

107@gws.ext.object.helper('upload') 

108class Object(gws.Node): 

109 """Upload helper, which receives file uploads in chunks and assembles them.""" 

110 

111 maxSize: int 

112 """Maximum upload size in bytes.""" 

113 maxChunkCount: int 

114 """Maximum number of chunks per upload.""" 

115 

116 def configure(self): 

117 self.maxSize = self.cfg('maxSize', default=1000) * 1024 * 1024 

118 self.maxChunkCount = max(1, self.maxSize // (500 * 1024)) # min. 500K chunks 

119 

120 def handle_chunk_request(self, req: gws.WebRequester, p: ChunkRequest) -> ChunkResponse: 

121 """Store a chunk of an upload, starting a new upload if ``uploadUid`` is empty. 

122 

123 Args: 

124 req: Web requester. 

125 p: Chunk request. 

126 

127 Returns: 

128 Response with the upload uid. 

129 

130 Raises: 

131 ``gws.BadRequestError``: If the chunk or the upload is invalid. 

132 """ 

133 try: 

134 up = self._save_chunk(p) 

135 return ChunkResponse(uploadUid=up.uid) 

136 except Error as exc: 

137 gws.log.exception() 

138 raise gws.BadRequestError('upload_error') from exc 

139 

140 def get_upload(self, uid: str) -> Upload: 

141 """Get a finished upload, joining its chunks into one file on the first call. 

142 

143 Args: 

144 uid: Upload uid. 

145 

146 Returns: 

147 The upload, with ``path`` pointing to the assembled file. 

148 

149 Raises: 

150 ``Error``: If the upload is not found, not complete or the file size does not match. 

151 """ 

152 up = self._load_upload(uid) 

153 out_path = _base_path(up.uid, 'out') 

154 

155 if not gws.u.is_file(out_path): 

156 with gws.u.server_lock(f'upload_{up.uid}'): 

157 self._finalize(up, out_path) 

158 

159 up.path = out_path 

160 return up 

161 

162 ## 

163 

164 def _save_chunk(self, p: ChunkRequest) -> Upload: 

165 """Validate a chunk and write it to the upload directory.""" 

166 up = self._load_upload(p.uploadUid) if p.uploadUid else self._create_upload(p) 

167 

168 if p.chunkNumber < 0 or p.chunkNumber >= up.chunkCount: 

169 raise Error(f'upload: {up.uid!r} invalid chunk number') 

170 

171 if len(p.content) > up.totalSize: 

172 raise Error(f'upload: {up.uid!r} invalid chunk size') 

173 

174 with gws.u.server_lock(f'upload_{up.uid}'): 

175 gws.u.write_file_b(_base_path(up.uid, p.chunkNumber), p.content) 

176 

177 return up 

178 

179 def _finalize(self, up: Upload, out_path): 

180 """Join all chunks into the output file and delete the chunks.""" 

181 chunks = [_base_path(up.uid, n) for n in range(0, up.chunkCount)] 

182 complete = all(gws.u.is_file(c) for c in chunks) 

183 if not complete: 

184 raise Error(f'upload: {up.uid!r}: incomplete') 

185 

186 tmp_path = out_path + '.tmp' 

187 with open(tmp_path, 'wb') as fp_all: 

188 for c in chunks: 

189 try: 

190 with open(c, 'rb') as fp: 

191 shutil.copyfileobj(fp, fp_all) 

192 except (OSError, IOError) as exc: 

193 raise Error(f'upload: {up.uid!r}: IO error') from exc 

194 

195 if gws.lib.osx.file_size(tmp_path) != up.totalSize: 

196 raise Error(f'upload: {up.uid!r}: invalid file size') 

197 

198 # @TODO check checksums as well? 

199 

200 try: 

201 gws.lib.osx.rename(tmp_path, out_path) 

202 except OSError: 

203 raise Error(f'upload: {up.uid!r}: move error') 

204 

205 for c in chunks: 

206 gws.lib.osx.unlink(c) 

207 

208 def _create_upload(self, p: ChunkRequest) -> Upload: 

209 """Validate the size and chunk count and create a new upload state.""" 

210 if p.totalSize <= 0 or p.totalSize > self.maxSize: 

211 raise Error(f'upload: invalid total size {p.totalSize!r}') 

212 if p.chunkCount <= 0 or p.chunkCount > self.maxChunkCount: 

213 raise Error(f'upload: invalid chunk count {p.chunkCount!r}') 

214 

215 uid = gws.u.random_string(64) 

216 up = Upload( 

217 uid=uid, 

218 fileName=p.fileName, 

219 totalSize=p.totalSize, 

220 chunkCount=p.chunkCount, 

221 path='', 

222 ) 

223 gws.lib.jsonx.to_path(_base_path(uid, 'state'), up) 

224 return up 

225 

226 def _load_upload(self, uid) -> Upload: 

227 """Load the state of an existing upload.""" 

228 if not uid.isalnum(): 

229 raise Error(f'upload: invalid uid {uid!r}') 

230 try: 

231 return Upload(gws.lib.jsonx.from_path(_base_path(uid, 'state'))) 

232 except gws.lib.jsonx.Error as exc: 

233 raise Error(f'upload: not found {uid!r}') from exc 

234 

235 

236def _base_path(uid, p): 

237 """Return the path of a file in the upload directory.""" 

238 return gws.u.ephemeral_dir(f'upload_{uid}') + f'/{p}'