Coverage for gramex\handlers\uploadhandler.py: 86%

# Coverage for gramex\handlers\uploadhandler.py : 86%

## 155 statements 135 run 20 missing 0 excluded 9 partial

Hot-keys on this page

r m x p toggle line displays

j k next/prev highlighted chunk

0 (zero) top of page

1 (one) first highlighted chunk

1 import os

2 import six

3 import time

4 import json

5 import gramex

6 import shutil

7 import mimetypes

8 import tornado.gen

9 from datetime import datetime

10 from six.moves import zip_longest

11 from orderedattrdict import AttrDict

12 from tornado.web import HTTPError

13 from gramex.config import app_log

14 from gramex.cache import HDF5Store, get_store

15 from gramex.transforms import build_transform

16 from gramex.http import FORBIDDEN, INTERNAL_SERVER_ERROR

17 from .basehandler import BaseHandler

18

19 MILLISECONDS = 1000

20

21

22 class FileUpload(object):

23 stores = {}

24

25 def __init__(self, path, keys=None, **kwargs):

26 if keys is None:

27 keys = {}

28 for cat in ('file', 'delete', 'save'):

29 keys.setdefault(cat, [cat])

30 if not isinstance(keys[cat], list):

31 if isinstance(keys[cat], six.string_types): 31 ↛ 34line 31 didn't jump to line 34, because the condition on line 31 was never false

32 keys[cat] = [keys[cat]]

33 else:

34 app_log.error('FileUpload: cat: %r must be a list or str', keys[cat])

35 self.keys = keys

36 self.path = os.path.abspath(path)

37 if not os.path.exists(self.path):

38 os.makedirs(self.path)

39

40 # store default: sqlite .meta.db

41 store_kwargs = kwargs.get('store', {

42 'type': 'sqlite',

43 'path': os.path.join(self.path, '.meta.db')

44 })

45 if self.path not in self.stores:

46 self.stores[self.path] = get_store(**store_kwargs)

47 self.store = self.stores[self.path]

48 old_store_path = os.path.abspath(os.path.join(self.path, '.meta.h5'))

49 store_path = os.path.abspath(getattr(self.store, 'path', None))

50 # migration: if type is not hdf5 but .meta.h5 exists, update store and remove

51 if (os.path.exists(old_store_path) and store_path != old_store_path): 51 ↛ 52line 51 didn't jump to line 52, because the condition on line 51 was never true

52 self._migrate_h5(old_store_path)

53

54 if 'file' not in keys: 54 ↛ 55line 54 didn't jump to line 55, because the condition on line 54 was never true

55 keys['file'] = ['file']

56 self.keys['file'] = keys['file'] if isinstance(keys['file'], list) else [keys['file']]

57

58 def _migrate_h5(self, old_store_path):

59 try:

60 old_store = HDF5Store(old_store_path, flush=5)

61 old_info = [(key, old_store.load(key)) for key in old_store.keys()]

62 for key, val in old_info:

63 self.store.dump(key, val)

64 self.store.flush()

65 old_store.close()

66 os.remove(old_store_path)

67 except Exception:

68 import sys

69 app_log.exception('FATAL: Cannot migrate: {}'.format(old_store_path))

70 sys.exit(1)

71

72 def info(self):

73 store = self.store

74 info = [(k, store.load(k)) for k in store.keys()]

75 return {k: v for k, v in info if v is not None}

76

77 def addfiles(self, handler):

78 filemetas = []

79 uploads = [upload for key in self.keys.get('file', [])

80 for upload in handler.request.files.get(key, [])]

81 filenames = [name for key in self.keys.get('save', [])

82 for name in handler.args.get(key, [])]

83 if_exists = getattr(handler, 'if_exists', 'unique')

84 for upload, filename in zip_longest(uploads, filenames, fillvalue=None):

85 filemeta = self.save_file(upload, filename, if_exists)

86 key = filemeta['file']

87 filemeta.update(

88 key=key,

89 user=handler.get_current_user(),

90 data=handler.args,

91 )

92 filemeta = handler.transforms(filemeta)

93 self.store.dump(key, filemeta)

94 filemetas.append(filemeta)

95 return filemetas

96

97 def save_file(self, upload, filename, if_exists):

98 original_name = upload.get('filename', None)

99 filemeta = AttrDict(filename=original_name)

100 filename = filename or original_name or 'data.bin'

101 filepath = os.path.join(self.path, filename)

102 # Security check: don't allow files to be written outside path:

103 if not os.path.realpath(filepath).startswith(os.path.realpath(self.path)):

104 raise HTTPError(FORBIDDEN, reason='FileUpload: filename %s is outside path: %s' % (

105 filename, self.path))

106 if os.path.exists(filepath):

107 if if_exists == 'error':

108 raise HTTPError(FORBIDDEN, reason='FileUpload: file exists: %s' % filename)

109 elif if_exists == 'unique':

110 # Rename to file.1.ext or file.2.ext etc -- whatever's available

111 name, ext = os.path.splitext(filepath)

112 name_pattern = name + '.%s' + ext

113 i = 1

114 while os.path.exists(name_pattern % i):

115 i += 1

116 filepath = name_pattern % i

117 elif if_exists == 'backup':

118 name, ext = os.path.splitext(filepath)

119 backup = '{}.{:%Y%m%d-%H%M%S}{}'.format(name, datetime.now(), ext)

120 shutil.copyfile(filepath, backup)

121 filemeta['backup'] = os.path.relpath(backup, self.path).replace(os.path.sep, '/')

122 elif if_exists != 'overwrite': 122 ↛ 123line 122 didn't jump to line 123, because the condition on line 122 was never true

123 raise HTTPError(INTERNAL_SERVER_ERROR,

124 reason='FileUpload: if_exists: %s invalid' % if_exists)

125 # Create the directory to write in, if reuqired

126 folder = os.path.dirname(filepath)

127 if not os.path.exists(folder):

128 os.makedirs(folder)

129 # Save the file

130 with open(filepath, 'wb') as handle:

131 handle.write(upload['body'])

132 mime = upload['content_type'] or mimetypes.guess_type(filepath, strict=False)[0]

133 filemeta.update(

134 file=os.path.relpath(filepath, self.path).replace(os.path.sep, '/'),

135 size=os.stat(filepath).st_size,

136 mime=mime or 'application/octet-stream',

137 created=time.time() * MILLISECONDS, # JS parseable timestamp

138 )

139 return filemeta

140

141 def deletefiles(self, handler):

142 status = []

143 for delete_key in self.keys.get('delete', []):

144 for key in handler.args.get(delete_key, []):

145 stat = {'success': False, 'key': key}

146 if key in self.store.keys(): 146 ↛ 152line 146 didn't jump to line 152, because the condition on line 146 was never false

147 path = os.path.join(self.path, key)

148 if os.path.exists(path): 148 ↛ 152line 148 didn't jump to line 152, because the condition on line 148 was never false

149 os.remove(path)

150 self.store.dump(key, None)

151 stat['success'] = True

152 status.append(stat)

153 return status

154

155

156 class UploadHandler(BaseHandler):

157 '''

158 UploadHandler lets users upload files. Here's a typical configuration::

159

160 path: /$GRAMEXDATA/apps/appname/ # Save files here

161 keys: [upload, file] # <input name=""> can be upload / file

162 store:

163 type: sqlite # Store metadata in a SQLite store

164 path: ... # ... at the specified path

165 redirect: # After uploading the file,

166 query: next # ... redirect to ?next=

167 url: /$YAMLURL/ # ... else to this directory

168 '''

169 @classmethod

170 def setup(cls, path, keys=None, if_exists='unique', transform=None, methods=[], **kwargs):

171 super(UploadHandler, cls).setup(**kwargs)

172 cls.if_exists = if_exists

173 # FileUpload uses the store= from **kwargs and ignores the rest

174 cls.uploader = FileUpload(path, keys=keys, **kwargs)

175

176 # methods=['get'] will show all file into as JSON on GET

177 if not isinstance(methods, list):

178 methods = [methods]

179 methods = {method.lower() for method in methods}

180 if 'get' in methods:

181 cls.get = cls.fileinfo

182

183 cls.transform = []

184 if transform is not None:

185 if isinstance(transform, dict) and 'function' in transform:

186 cls.transform.append(build_transform(

187 transform, vars=AttrDict((('content', None), ('handler', None))),

188 filename='url:%s' % cls.name))

189 else:

190 app_log.error('UploadHandler %s: no function: in transform: %r',

191 cls.name, transform)

192

193 @tornado.gen.coroutine

194 def fileinfo(self, *args, **kwargs):

195 self.set_header('Content-Type', 'application/json')

196 self.write(json.dumps(self.uploader.info(), indent=2))

197

198 @tornado.gen.coroutine

199 def post(self, *args, **kwargs):

200 if self.redirects: 200 ↛ 201line 200 didn't jump to line 201, because the condition on line 200 was never true

201 self.save_redirect_page()

202 upload = yield gramex.service.threadpool.submit(self.uploader.addfiles, self)

203 delete = yield gramex.service.threadpool.submit(self.uploader.deletefiles, self)

204 self.set_header('Content-Type', 'application/json')

205 self.write(json.dumps({'upload': upload, 'delete': delete},

206 ensure_ascii=True, separators=(',', ':')))

207 if self.redirects: 207 ↛ 208line 207 didn't jump to line 208, because the condition on line 207 was never true

208 self.redirect_next()

209

210 def transforms(self, content):

211 for transform in self.transform:

212 for value in transform(content, self):

213 if isinstance(value, dict):

214 content = value

215 elif value is not None: 215 ↛ 212line 215 didn't jump to line 212, because the condition on line 215 was never false

216 app_log.error('UploadHandler %s: transform returned %r, not dict',

217 self.name, value)

218 return content