1import json
2import mimetypes
3import os
4from base64 import decodebytes, encodebytes
5from datetime import datetime
6from io import BytesIO
7from urllib.parse import unquote
8
9import nbformat
10from jupyter_server.services.contents.manager import ContentsManager
11from tornado import web
12from traitlets import default
13
14from omegaml.notebook.checkpoints import NoOpCheckpoints
15
16
[docs]
17class OmegaStoreContentsManager(ContentsManager):
18 """
19 Jupyter notebook storage manager for omegaml
20
21 Adopted from notebook/services/contents/filemanager.py
22
23 This requires a properly configured omegaml instance.
24 see https://jupyter-server.readthedocs.io/en/latest/developers/contents.html
25 """
26
27 def __init__(self, **kwargs):
28 # pass omega= for testing purpose
29 self._omega = kwargs.pop('omega', None)
30 super().__init__(**kwargs)
31
32 def _checkpoints_class_default(self):
33 return NoOpCheckpoints
34
35 @property
36 def omega(self):
37 """
38 return the omega instance used by the contents manager
39 """
40 if self._omega is None:
41 import omegaml as om
42
43 self._omega = om
44 self._omega.jobs._include_dir_placeholder = True
45 return self._omega
46
47 @property
48 def store(self):
49 """
50 return the OmageStore for jobs (notebooks)
51 """
52 return self.omega.jobs
53
54 @property
55 def _dir_placeholder(self):
56 return self.omega.jobs._dir_placeholder
57
58 def get(self, path, content=True, type=None, format=None):
59 """
60 get an entry in the store
61
62 this is called by the contents engine to get the contents of the jobs
63 store.
64 """
65 path = unquote(path).strip('/')
66 if type == 'notebook' or (type is None and path.endswith('.ipynb')):
67 model = self._notebook_model(path, content=content)
68 elif type == 'file':
69 model = self._file_model(path, content=content)
70 elif type in (None, 'directory'):
71 # jupyterlab passes None to get directory
72 # we never return content with a directory listing to save time
73 # the frontend will request the specific contents
74 model = self._dir_model(path, content=False)
75 else:
76 raise web.HTTPError(404, f'Type {type} at {path} is not supported')
77 return model
78
79 def save(self, model, path):
80 """
81 save an entry in the store
82
83 this is called by the contents engine to store a notebook
84 """
85 om = self.omega
86 path = unquote(path).strip('/')
87 type = model.get('type')
88 name = model.get('name')
89 self.run_pre_save_hooks(model=model, path=path)
90 if type is None:
91 raise web.HTTPError(400, 'No file type provided')
92 try:
93 if type == 'notebook' or (type == 'file' and path.endswith('.ipynb')):
94 content = model.get('content')
95 # parse input
96 if model.get('format') == 'text' and isinstance(content, bytes):
97 content = content.decode('utf8')
98 elif model.get('format') == 'base64':
99 if isinstance(content, str):
100 content = content.encode('ascii')
101 content = decodebytes(content).decode('utf8')
102 while isinstance(content, str):
103 content = json.loads(content)
104 model['content'] = content
105 model['format'] = None
106 if content is None or not isinstance(content, dict):
107 raise web.HTTPError(400, 'No file content provided or wrong format')
108 # create notebook
109 nb = nbformat.from_dict(content)
110 self.check_and_sign(nb, path)
111 self.omega.jobs.put(nb, path)
112 self.validate_notebook_model(model)
113 validation_message = model.get('message', None)
114 model = self.get(path, content=False, type=type)
115 if validation_message:
116 model['message'] = validation_message
117 elif type == 'directory':
118 ph_name = '{path}/{self._dir_placeholder}'.format(**locals()).strip('/')
119 self.omega.jobs.create("#placeholder", ph_name)
120 model = self.get(path, content=False, type=type)
121 model['content'] = None
122 model['format'] = None
123 validation_message = None
124 elif type == 'file':
125 content = model.get('content')
126 fmt = model.get('format')
127 self._save_file(path, content, fmt)
128 model = self.get(path, content=False, type=type)
129 else:
130 raise web.HTTPError(400, "Unhandled contents type: %s" % model['type'])
131 except web.HTTPError:
132 raise
133 except Exception as e:
134 self.log.error('Error while saving file: %s %s', path, e, exc_info=True)
135 raise web.HTTPError(500, 'Unexpected error while saving file: %s %s' % (path, e))
136 return model
137
138 def delete_file(self, path):
139 """
140 delete an entry
141
142 this is called by the contents engine to delete an entry
143 """
144 path = unquote(path).strip('/')
145 try:
146 self.omega.jobs.drop(path)
147 except Exception:
148 self.omega.jobs.drop(path + '/' + self._dir_placeholder)
149
150 def rename_file(self, old_path, new_path):
151 """
152 rename a file
153
154 this is called by the contents engine to rename an entry
155 """
156 old_path = unquote(old_path).strip('/')
157 new_path = unquote(new_path).strip('/')
158 # check file or directory
159 if self.file_exists(new_path):
160 raise web.HTTPError(409, 'Notebook already exists: %s' % new_path)
161 elif self.dir_exists(new_path):
162 raise web.HTTPError(409, 'Directory already exists: %s' % new_path)
163 # do the renaming
164 if self.dir_exists(old_path):
165 old_dirname = old_path + '/' + self._dir_placeholder
166 new_dirname = new_path + '/' + self._dir_placeholder
167 meta = self.omega.jobs.metadata(old_dirname)
168 meta.name = new_dirname
169 elif self.file_exists(old_path):
170 meta = self.omega.jobs.metadata(old_path)
171 meta.name = new_path
172 # rename on metadata. Note the gridfile instance stays the same
173 meta.save()
174
175 def exists(self, path):
176 """
177 Does a file or dir exist at the given collection in gridFS?
178 We do not have dir so dir_exists returns true.
179
180 :param path: (str) The relative path to the file's directory
181 (with '/' as separator)
182 :returns exists: (boo) The relative path to the file's directory (with '/' as separator)
183 """
184 path = unquote(path).strip('/')
185 return self.file_exists(path) or self.dir_exists(path)
186
187 def dir_exists(self, path=''):
188 """check if directory exists
189
190 Args:
191 path: name of directory
192
193 Returns:
194 True if directory exists
195 """
196 path = unquote(path).strip('/')
197 if path == '':
198 return True
199 pattern = rf'^{path}.*/({self._dir_placeholder}|.+)'
200 return len(self.omega.jobs.list(regexp=pattern)) > 0
201
202 def file_exists(self, path=""):
203 """check if file exists
204
205 Args:
206 path: name of file
207
208 Returns:
209 True if file exists
210 """
211 path = unquote(path).strip('/')
212 if not path:
213 return False
214 # always check for an actual file, not some sub path
215 pattern = rf'^{path}$'
216 does_exist = len(self.omega.jobs.list(regexp=pattern)) > 0
217 does_exist |= len(self.omega.datasets.list(regexp=pattern)) > 0
218 return does_exist
219
220 def is_hidden(self, path):
221 """check if path or file is hidden
222
223 Args:
224 path: name of file or path
225
226 Returns:
227 False, currently always returns false
228 """
229 return False
230
231 def _read_notebook(self, path, as_version=None):
232 path = unquote(path).strip('/')
233 return self.omega.jobs.get(path)
234
235 def _notebook_model(self, path, content=True, meta=None):
236 """
237 Build a notebook model
238 if content is requested, the notebook content will be populated
239 as a JSON structure (not double-serialized)
240 """
241 path = unquote(path).strip('/')
242 model = self._base_model(path)
243 model['type'] = 'notebook'
244 # always add accurate created and modified
245 meta = meta or self.omega.jobs.metadata(path)
246 if meta is not None:
247 model['created'] = meta.created
248 model['last_modified'] = meta.modified
249 if content:
250 nb = self._read_notebook(path, as_version=4)
251 if nb is None:
252 raise web.HTTPError(400, f"Cannot read non-file {path}")
253 self.mark_trusted_cells(nb, path)
254 model['content'] = nb
255 model['format'] = 'json'
256 self.validate_notebook_model(model)
257 return model
258
259 def _base_model(self, path, kind=None):
260 """Build the common base of a contents model"""
261 # http://jupyter-notebook.readthedocs.io/en/stable/extending/contents.html
262 path = unquote(path).strip('/')
263 last_modified = datetime.utcnow()
264 created = last_modified
265 # Create the base model.
266 model = {}
267 model['name'] = os.path.basename(path)
268 model['path'] = path
269 model['last_modified'] = last_modified
270 model['created'] = created
271 model['content'] = None
272 model['format'] = None
273 model['mimetype'] = None
274 model['writable'] = True
275 if kind:
276 model['type'] = kind
277 model['content'] = [] if kind == 'directory' else None
278 return model
279
280 def _dir_model(self, path, content=True):
281 """
282 Build a model to return all of the files in gridfs
283 if content is requested, will include a listing of the directory
284 """
285 # this looks like a seemingly simple task, it's carefully crafted
286 path = unquote(path).strip('/')
287 model = self._base_model(path, kind='directory')
288 model['format'] = 'json'
289 contents = model['content']
290 # get existing entries from a pattern that matches either
291 # top-level files: ([\w -]+\.[\w -]*)
292 # directories (files in): ([\w ]+/([\w ]+\.[\w]*))
293 # does not work:
294 # pattern = r'([\w ]+/_placeholder\.[\w]*)|([\w ]+\.[\w]*)$'
295 # it is too restrictive as entries can be generated without a placeholder
296 # so we get all, which may include sub/sub/file
297 # and we need to include them because we need to find sub/sub directories
298 # note \w is any word character (letter, digit, underscore)
299 # \s is any white space
300 # \d is any digit
301 # :_ match literally
302 # [^\/] matches any character except /
303 # pattern = r'([\w\s\-.\d:()+]+\/)?([\w\s\-.\d:()+]+\.[\w]*)$'
304 pattern = r'([^\/]+\/)?([^\/]+\.[^\/]*)$'
305 # if we're looking in an existing directory, prepend that
306 if path:
307 pattern = rf'{path}/{pattern}'
308 pattern = rf'^{pattern}'
309 entries = self.omega.jobs.list(regexp=pattern, raw=True, hidden=True, include_temp=True)
310 if path and not entries:
311 raise web.HTTPError(400, f"Directory not found {path}")
312 # by default assume the current path is listed already
313 directories = [path]
314 for meta in entries:
315 # get path of entry, e.g. sub/foo.ipynb => sub
316 entry_path = os.path.dirname(meta.name)
317 # if not part of listed directories yet, include
318 if entry_path not in directories:
319 entry = self._base_model(entry_path, kind='directory')
320 contents.append(entry)
321 directories.append(entry_path)
322 # ignore placeholder files
323 if meta.name.endswith(self._dir_placeholder):
324 continue
325 # only include files that are in the path we're listing
326 if entry_path != path:
327 continue
328 # include the actual file
329 try:
330 entry = self._notebook_model(meta.name, content=content, meta=meta)
331 except Exception as e:
332 msg = f'_dir_model error, cannot get {meta.name}, removing from list, exception {e!s}'
333 self.log.warning(msg)
334 else:
335 contents.append(entry)
336 return model
337
338 def _file_model(self, path, content=True, format=None):
339 model = self._base_model(path)
340 model['type'] = 'file'
341
342 model['mimetype'] = mimetypes.guess_type(path)[0]
343
344 if content:
345 content, format = self._read_file(path, format)
346 if model['mimetype'] is None:
347 default_mime = {
348 'text': 'text/plain',
349 'base64': 'application/octet-stream',
350 }[format]
351 model['mimetype'] = default_mime
352
353 model.update(
354 content=content,
355 format=format,
356 )
357 return model
358
359 def _read_file(self, os_path, format):
360 """Read a non-notebook file.
361
362 os_path: The path to be read.
363 format:
364 If 'text', the contents will be decoded as UTF-8.
365 If 'base64', the raw bytes contents will be encoded as base64.
366 If not specified, try to decode as UTF-8, and fall back to base64
367 """
368 if os_path.endswith('.ipynb'):
369 meta = self.omega.jobs.metadata(os_path)
370 else:
371 meta = self.omega.datasets.metadata(os_path)
372 if meta is None or meta.gridfile is None:
373 raise web.HTTPError(400, "Cannot read non-file %s" % os_path)
374
375 if meta.gridfile:
376 bcontent = meta.gridfile.read()
377 meta.gridfile.close()
378
379 if format is None or format == 'text':
380 # Try to interpret as unicode if format is unknown or if unicode
381 # was explicitly requested.
382 try:
383 return bcontent.decode('utf8'), 'text'
384 except UnicodeError:
385 if format == 'text':
386 raise web.HTTPError(400, "%s is not UTF-8 encoded" % os_path, reason='bad format')
387 return encodebytes(bcontent).decode('ascii'), 'base64'
388
389 def _save_file(self, os_path, content, format):
390 """Save content of a generic file."""
391 if format not in {'text', 'base64'}:
392 raise web.HTTPError(400, "Must specify format of file contents as 'text' or 'base64'")
393 try:
394 if format == 'text':
395 bcontent = content.encode('utf8')
396 else:
397 b64_bytes = content.encode('ascii')
398 bcontent = decodebytes(b64_bytes)
399 except Exception as e:
400 raise web.HTTPError(400, 'Encoding error saving %s: %s' % (os_path, e))
401
402 self.omega.datasets.put(BytesIO(bcontent), os_path)
403
404 @default('files_handler_params')
405 def _files_handler_params_default(self):
406 # avoid exception
407 # TypeError: StaticFileHandler.initialize() missing 1 required positional argument: 'path'
408 # issue: https://github.com/jupyter-server/jupyter_server/issues/1313
409 # What this does
410 # - ensure that ContentsManager.files_handler_class=FilesHandler gets has a path attribute on initialize()
411 # - this only happens when dealing with local files, e.g. when downloading a notebook
412 # - source of the issue is that ContentsManager.files_handler_params is not set
413 return {'path': self.root_dir}