Merge pull request #121 from dataiku/feature/dss90-jupyter-notebooks-public-api

Python wrapper for Jupyter Notebooks Public API

Author: william-lukusa

dataikuapi/dss/notebook.py

@@ -4,49 +4,82 @@ class DSSNotebook(object):
     """
     A Python/R/Scala notebook on the DSS instance
     """
-    def __init__(self, client, project_key, notebook_name, state=None):
+    def __init__(self, client, project_key, notebook_name, state=None, content=None):
        self.client = client
        self.project_key = project_key
        self.notebook_name = notebook_name
        self.state = state
+       self.content = content
        self.state_is_peek = True
 
     def unload(self, session_id=None):
         """
-        Stop the notebook and release its resources
+        Stop this Jupyter notebook and release its resources
         """
-        state = self.get_state()
-        if state is None:
+        sessions = self.get_sessions()
+        if sessions is None:
             raise Exception("Notebook isn't running")
-        if state.get('activeSessions', None) is None:
-            raise Exception("Notebook isn't running")
-        if len(state['activeSessions']) == 0:
+        if len(sessions) == 0:
             raise Exception("Notebook isn't running")
         if session_id is None:
-            if len(state['activeSessions']) > 1:
+            if len(sessions) > 1:
                 raise Exception("Several sessions of the notebook are running, choose one")
             else:
-                session_id = state['activeSessions'][0].get('sessionId', None)
-        return self.client._perform_json("DELETE", "/projects/%s/notebooks/" % self.project_key, params={'notebookName' : self.notebook_name, 'sessionId' : session_id})
+                session_id = sessions[0].get('sessionId', None)
+        return self.client._perform_json("DELETE",
+                                         "/projects/%s/jupyter-notebooks/%s/sessions/%s" % (self.project_key, self.notebook_name, session_id))
 
-    def get_state(self):
+    def get_state(self, refresh=False):
         """
-        Get the status of the notebook
+        Get the status of this Jupyter notebook
+
+        :param bool refresh: if True, get the status of the notebook from the backend
         """
-        if self.state is None:
-            self.state = self.client._perform_json("GET", "/projects/%s/notebooks/" % self.project_key, params={'notebookName' : self.notebook_name})
+        notebook_states = self.client._perform_json("GET",
+                                                    "/projects/%s/jupyter-notebooks/" % self.project_key,
+                                                    params={"active": False})
+        if self.state is None or refresh:
+            for notebook_state in notebook_states:
+                if notebook_state.get("name") == self.notebook_name:
+                    self.state = notebook_state
+                    break
         return self.state
 
     def get_sessions(self):
         """
-        Get the list of the running sessions of this notebook
+        Get the list of running sessions of this Jupyter notebook
         """
-        state = self.get_state()
-        if state is None:
-            raise Exception("Notebook isn't running")
-        if state.get('activeSessions', None) is None:
-            raise Exception("Notebook isn't running")
-        return state['activeSessions']
+
+        if self.state is None:
+            self.state = {}
+        sessions = self.client._perform_json("GET",
+                                             "/projects/%s/jupyter-notebooks/%s/sessions" % (self.project_key, self.notebook_name))
+        self.state["activeSessions"] = sessions
+        return sessions
+
+    def get_content(self):
+        """
+        Get the content of this Jupyter notebook (metadata, cells, nbformat)
+        """
+        if self.content is None:
+            self.content = self.client._perform_json("GET",
+                                                     "/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name))
+        return self.content
+
+    def save(self):
+        """
+        Save the content of this Jupyter notebook
+        """
+        return self.client._perform_json("PUT",
+                                         "/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name),
+                                         body=self.content)
+
+    def delete(self):
+        """
+        Delete this Jupyter notebook and stop all of its active sessions.
+        """
+        return self.client._perform_json("DELETE",
+                                         "/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name))
 
     ########################################################
     # Discussions

dataikuapi/dss/project.py

@@ -825,6 +825,57 @@ def new_job_definition_builder(self, job_type='NON_RECURSIVE_FORCED_BUILD'):
         warnings.warn("new_job_definition_builder is deprecated, please use new_job", DeprecationWarning)
         return JobDefinitionBuilder(self, job_type)
 
+    ########################################################
+    # Jupyter Notebooks
+    ########################################################
+    
+    def list_jupyter_notebooks(self, as_objects=True, active=False):
+        """
+        List the jupyter notebooks of a project.
+
+        :param bool as_objects: if True, return the jupyter notebooks as a :class:`dataikuapi.dss.notebook.DSSNotebook`
+                        notebook handles instead of raw JSON
+        :param bool active: if True, only return currently running jupyter notebooks.
+
+
+        :returns: The list of the notebooks - see as_objects for more information
+        :rtype: list
+        """
+        notebooks = self.client._perform_json("GET", "/projects/%s/jupyter-notebooks/" % self.project_key,
+                                         params={"active": active})
+        if as_objects:
+            return [DSSNotebook(self.client, notebook_state['projectKey'], notebook_state['name'], state=notebook_state) for notebook_state in notebooks]
+        else:
+            return notebooks
+
+    def get_jupyter_notebook(self, notebook_name):
+        """
+        Get a handle to interact with a specific jupyter notebook
+
+        :param str notebook_name: The name of the jupyter notebook to retrieve
+        :returns: A handle to interact with this jupyter notebook
+        :rtype: :class:`~dataikuapi.dss.notebook.DSSNotebook` jupyter notebook handle
+        """
+        notebook_content = self.client._perform_json("GET",
+                                             "/projects/%s/jupyter-notebooks/%s" % (self.project_key, notebook_name))
+        return DSSNotebook(self.client, self.project_key, notebook_name, content=notebook_content)
+
+    def create_jupyter_notebook(self, notebook_name, notebook_content):
+        """
+        Create a new jupyter notebook and get a handle to interact with it
+
+        :param str notebook_name: the name of the notebook to create
+        :param dict notebook_content: the data of the notebook to create, as a dict.
+            The data will be converted to a JSON string internally.
+            Use ``get_content()`` on a similar existing ``DSSNotebook`` object in order to get a sample definition object
+        :returns: A handle to interact with the newly created jupyter notebook
+        :rtype: :class:`~dataikuapi.dss.notebook.DSSNotebook` jupyter notebook handle
+        """
+        created_notebook_content = self.client._perform_json("POST",
+                                  "/projects/%s/jupyter-notebooks/%s" % (self.project_key, notebook_name),
+                                  body=notebook_content)
+        return DSSNotebook(self.client, self.project_key, notebook_name, content=created_notebook_content)
+
     ########################################################
     # Continuous activities
     ########################################################