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

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 @gws.ext.command.api('myUpload') 

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

9 # check permissions, etc... 

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

11 return helper.handle_chunk_request(req, p) 

12 ... 

13 

14The 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. 

15 

16Once 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. 

17 

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""" 

31import shutil 

32 

33import gws 

34import gws.lib.jsonx 

35import gws.lib.osx 

36 

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

38 

39 

40class ChunkRequest(gws.Request): 

41 uploadUid: str = '' 

42 fileName: str 

43 totalSize: int 

44 chunkNumber: int 

45 chunkCount: int 

46 content: bytes 

47 

48 

49class ChunkResponse(gws.Response): 

50 uploadUid: str 

51 

52 

53class Upload(gws.Data): 

54 uploadUid: str 

55 fileName: str 

56 totalSize: int 

57 path: str 

58 

59 

60class _Params(gws.Data): 

61 uid: str 

62 fileName: str 

63 totalSize: int 

64 chunkCount: int 

65 

66 

67class Error(gws.Error): 

68 pass 

69 

70 

71class Object(gws.Node): 

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

73 try: 

74 ps = self._save_chunk(p) 

75 return ChunkResponse(uploadUid=ps.uid) 

76 except Error as exc: 

77 gws.log.exception() 

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

79 

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

81 ps = self._get_params(uid) 

82 dd = self._base_dir(ps.uid) 

83 out_path = f'{dd}/out' 

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

85 with gws.u.server_lock(f'upload_{ps.uid}'): 

86 self._finalize(ps, out_path) 

87 return Upload(uploadUid=ps.uid, path=out_path, fileName=ps.fileName) 

88 

89 ## 

90 

91 def _save_chunk(self, p: ChunkRequest) -> _Params: 

92 ps = self._get_params(p.uploadUid) if p.uploadUid else self._create_upload(p) 

93 

94 if p.chunkNumber < 0 or p.chunkNumber >= ps.chunkCount: 

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

96 

97 dd = self._base_dir(ps.uid) 

98 

99 with gws.u.server_lock(f'upload_{ps.uid}'): 

100 gws.u.write_file_b(f'{dd}/{p.chunkNumber}', p.content) 

101 

102 return ps 

103 

104 def _get_chunks(self, ps: _Params): 

105 dd = self._base_dir(ps.uid) 

106 chunks = [f'{dd}/{n}' for n in range(0, ps.chunkCount)] 

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

108 return chunks, complete 

109 

110 def _finalize(self, ps: _Params, out_path): 

111 chunks, complete = self._get_chunks(ps) 

112 if not complete: 

113 raise Error(f'upload: {ps.uid!r} incomplete') 

114 

115 tmp_path = out_path + '.tmp' 

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

117 for c in chunks: 

118 try: 

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

120 shutil.copyfileobj(fp, fp_all) 

121 except (OSError, IOError) as exc: 

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

123 

124 if gws.lib.osx.file_size(tmp_path) != ps.totalSize: 

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

126 

127 # @TODO check checksums as well? 

128 

129 try: 

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

131 except OSError: 

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

133 

134 for c in chunks: 

135 gws.lib.osx.unlink(c) 

136 

137 def _create_upload(self, p: ChunkRequest) -> _Params: 

138 uid = gws.u.random_string(64) 

139 dd = self._base_dir(uid) 

140 

141 gws.lib.jsonx.to_path(f'{dd}/s.json', _Params( 

142 uid=uid, 

143 fileName=p.fileName, 

144 totalSize=p.totalSize, 

145 chunkCount=p.chunkCount, 

146 )) 

147 

148 return self._get_params(uid) 

149 

150 def _get_params(self, uid): 

151 if not uid.isalnum(): 

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

153 

154 dd = self._base_dir(uid) 

155 

156 try: 

157 return _Params(gws.lib.jsonx.from_path(f'{dd}/s.json')) 

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

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

160 

161 def _base_dir(self, uid): 

162 return gws.u.ephemeral_dir(f'upload_{uid}')