largefilemanager.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. from notebook.services.contents.filemanager import FileContentsManager
  2. from contextlib import contextmanager
  3. from tornado import web
  4. import nbformat
  5. import base64
  6. import os, io
  7. class LargeFileManager(FileContentsManager):
  8. """Handle large file upload."""
  9. def save(self, model, path=''):
  10. """Save the file model and return the model with no content."""
  11. chunk = model.get('chunk', None)
  12. if chunk is not None:
  13. path = path.strip('/')
  14. if 'type' not in model:
  15. raise web.HTTPError(400, u'No file type provided')
  16. if model['type'] != 'file':
  17. raise web.HTTPError(400, u'File type "{}" is not supported for large file transfer'.format(model['type']))
  18. if 'content' not in model and model['type'] != 'directory':
  19. raise web.HTTPError(400, u'No file content provided')
  20. os_path = self._get_os_path(path)
  21. try:
  22. if chunk == 1:
  23. self.log.debug("Saving %s", os_path)
  24. self.run_pre_save_hook(model=model, path=path)
  25. super(LargeFileManager, self)._save_file(os_path, model['content'], model.get('format'))
  26. else:
  27. self._save_large_file(os_path, model['content'], model.get('format'))
  28. except web.HTTPError:
  29. raise
  30. except Exception as e:
  31. self.log.error(u'Error while saving file: %s %s', path, e, exc_info=True)
  32. raise web.HTTPError(500, u'Unexpected error while saving file: %s %s' % (path, e))
  33. model = self.get(path, content=False)
  34. # Last chunk
  35. if chunk == -1:
  36. self.run_post_save_hook(model=model, os_path=os_path)
  37. return model
  38. else:
  39. return super(LargeFileManager, self).save(model, path)
  40. def _save_large_file(self, os_path, content, format):
  41. """Save content of a generic file."""
  42. if format not in {'text', 'base64'}:
  43. raise web.HTTPError(
  44. 400,
  45. "Must specify format of file contents as 'text' or 'base64'",
  46. )
  47. try:
  48. if format == 'text':
  49. bcontent = content.encode('utf8')
  50. else:
  51. b64_bytes = content.encode('ascii')
  52. bcontent = base64.b64decode(b64_bytes)
  53. except Exception as e:
  54. raise web.HTTPError(
  55. 400, u'Encoding error saving %s: %s' % (os_path, e)
  56. )
  57. with self.perm_to_403(os_path):
  58. if os.path.islink(os_path):
  59. os_path = os.path.join(os.path.dirname(os_path), os.readlink(os_path))
  60. with io.open(os_path, 'ab') as f:
  61. f.write(bcontent)