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

1"""Manage chunked uploads. 

2 

3In your action, declare an endpoint with ``p: ChunkRequest`` as a parameter. This endpoint should invoke ``handle_chunk_request``:: 

4 

5 import gws.plugin.upload_helper as uh 

6 

7 

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

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

10 # check permissions, etc... 

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

12 return helper.handle_chunk_request(req, p) 

13 ... 

14 

15The client sends chunks to this endpoint, one by one. Each chunk contains the file name and total size. The first chunk has an empty ``uploadUid``, indicating a new upload. Subsequent chunks must provide a valid ``uploadUid``. The handler responds with an ``uploadUid``. Each chunk must have a serial number, starting from 0. Chunks can come in any order. 

16 

17Once the client decides that the upload is complete, it proceeds with invoking some other endpoint of your action, mentioning the ``uploadUid`` returned by the first chunk. The endpoint should invoke ``get_upload`` to retrieve the final file. The file is stored in a temporary location and should be moved to a permanent location if necessary:: 

18 

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

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

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

22 try: 

23 upload = helper.get_upload(p.uploadUid) 

24 except uh.Error: 

25 ...upload not ready yet... 

26 ...process(upload.path) 

27 

28 

29 

30""" 

31 

32import shutil 

33 

34import gws 

35import gws.lib.jsonx 

36import gws.lib.osx 

37 

38gws.ext.new.helper('upload') 

39 

40 

41class Config(gws.Config): 

42 """Upload helper.""" 

43 

44 maxSize: int = 1000 

45 """Maximum upload size in megabytes.""" 

46 

47 

48class ChunkRequest(gws.Request): 

49 uploadUid: str = '' 

50 fileName: str 

51 totalSize: int 

52 chunkNumber: int 

53 chunkCount: int 

54 content: bytes 

55 

56 

57class ChunkResponse(gws.Response): 

58 uploadUid: str 

59 

60 

61class Upload(gws.Data): 

62 uid: str 

63 fileName: str 

64 totalSize: int 

65 chunkCount: int 

66 path: str 

67 

68 

69class Error(gws.Error): 

70 pass 

71 

72 

73class Object(gws.Node): 

74 maxSize: int 

75 maxChunkCount: int 

76 

77 def configure(self): 

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

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

80 

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

82 try: 

83 up = self._save_chunk(p) 

84 return ChunkResponse(uploadUid=up.uid) 

85 except Error as exc: 

86 gws.log.exception() 

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

88 

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

90 up = self._load_upload(uid) 

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

92 

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

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

95 self._finalize(up, out_path) 

96 

97 up.path = out_path 

98 return up 

99 

100 ## 

101 

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

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

104 

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

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

107 

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

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

110 

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

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

113 

114 return up 

115 

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

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

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

119 if not complete: 

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

121 

122 tmp_path = out_path + '.tmp' 

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

124 for c in chunks: 

125 try: 

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

127 shutil.copyfileobj(fp, fp_all) 

128 except (OSError, IOError) as exc: 

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

130 

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

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

133 

134 # @TODO check checksums as well? 

135 

136 try: 

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

138 except OSError: 

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

140 

141 for c in chunks: 

142 gws.lib.osx.unlink(c) 

143 

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

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

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

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

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

149 

150 uid = gws.u.random_string(64) 

151 up = Upload( 

152 uid=uid, 

153 fileName=p.fileName, 

154 totalSize=p.totalSize, 

155 chunkCount=p.chunkCount, 

156 path='', 

157 ) 

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

159 return up 

160 

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

162 if not uid.isalnum(): 

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

164 try: 

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

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

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

168 

169 

170def _base_path(uid, p): 

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