server.py

from notebook.utils import url_path_join
from notebook.base.handlers import IPythonHandler
from pybis import Openbis
import numpy as np

import os
from urllib.parse import unquote
import yaml

openbis_connections = {}


def _jupyter_server_extension_paths():
    return [{'module': 'jupyter-openbis-extension.server'}]

def _load_configuration(paths, filename='openbis-connections.yaml'):

    if paths is None:
        paths = []
        home = os.path.expanduser("~")
        paths.append(os.path.join(home, '.jupyter'))

    # look in all config file paths of jupyter
    # for openbis connection files and load them
    connections = []
    for path in paths:
        abs_filename = os.path.join(path, filename)
        if os.path.isfile(abs_filename):
            with open(abs_filename, 'r') as stream:
                try:
                    config = yaml.safe_load(stream)
                    for connection in config['connections']:
                        connections.append(connection)
                except yaml.YAMLexception as exc:
                    print(exc)
                    return None

    return connections
        

def load_jupyter_server_extension(nb_server_app):
    """Call when the extension is loaded.
    :param nb_server_app: Handle to the Notebook webserver instance.
    """

    # load the configuration file
    # and register the openBIS connections.
    # If username and password is available, try to connect to the server
    connections = _load_configuration(
        paths    = nb_server_app.config_file_paths,
        filename = 'openbis-connections.yaml'
    )

    for connection_info in connections:
        conn = register_connection(connection_info)
        print("Registered: {}".format(conn.url))
        if conn.username and conn.password:
            try:
                conn.login()
                print("Successfully connected to: {}".format(conn.url))
            except ValueError:
                print("Incorrect username or password for: {}".format(conn.url))
            except Exception:
                print("Cannot establish connection to: {}".format(conn.url))

    # Add URL handlers to our web_app
    # see Tornado documentation: https://www.tornadoweb.org
    web_app = nb_server_app.web_app
    host_pattern = '.*$'
    base_url = web_app.settings['base_url']

    # get the file list
    web_app.add_handlers(
        host_pattern,
        [(url_path_join( base_url, '/general/filelist'),
            FileListHandler
        )]
    )


    # DataSet download
    web_app.add_handlers(
        host_pattern, 
        [(url_path_join(
            base_url,
            '/openbis/dataset/(?P<connection_name>.*)?/(?P<permId>.*)?/(?P<downloadPath>.*)'), 
            DataSetDownloadHandler
        )]
    )

    # DataSet upload
    web_app.add_handlers( host_pattern, [(
            url_path_join( base_url, '/openbis/dataset/(?P<connection_name>.*)' ),
            DataSetUploadHandler
        )]
    )

    # DataSet-Types
    web_app.add_handlers(
        host_pattern,
        [(
            url_path_join( base_url, '/openbis/datasetTypes/(?P<connection_name>.*)'
            ),
            DataSetTypesHandler
        )]
    )

    # DataSets for Sample identifier/permId
    web_app.add_handlers(
        host_pattern, 
        [(  
            url_path_join(
                base_url,
                '/openbis/sample/(?P<connection_name>.*)?/(?P<permId>.*)'
            ), 
            SampleHandler
        )]
    )

    # OpenBIS connections
    web_app.add_handlers(
        host_pattern,
        [(
            url_path_join(
                base_url,
                '/openbis/conns'
            ),
            OpenBISConnections
        )]
    )

    # Modify / reconnect to a connection
    web_app.add_handlers(
        host_pattern, 
        [(
            url_path_join(
                    base_url,
                    '/openbis/conn/(?P<connection_name>.*)'
            ),
            OpenBISConnectionHandler
        )]
    )

    print("pybis loaded: {}".format(Openbis))


def register_connection(connection_info):

    conn = OpenBISConnection(
        name                = connection_info.get('name'),
        url                 = connection_info.get('url'),
        verify_certificates = connection_info.get('verify_certificates', False),
        username            = connection_info.get('username'),
        password            = connection_info.get('password'),
        status              = 'not connected',
    )
    openbis_connections[conn.name] = conn
    return conn


class OpenBISConnection:
    """register an openBIS connection
    """

    def __init__(self, **kwargs):
        for needed_key in ['name', 'url']:
            if needed_key not in kwargs:
                raise KeyError("{} is missing".format(needed_key))

        for key in kwargs:
            setattr(self, key, kwargs[key])

        openbis = Openbis(
            url = self.url,
            verify_certificates = self.verify_certificates
        )
        self.openbis = openbis
        self.status = "not connected"

    def is_session_active(self):
        return self.openbis.is_session_active()

    def check_status(self):
        if self.openbis.is_session_active():
            self.status = "connected"
        else:
            self.status = "not connected"

    def login(self, username=None, password=None):
        if username is None:
            username=self.username
        if password is None:
            password=self.password
        self.openbis.login(
            username = username,
            password = password
        )
        # store username and password in memory
        self.username = username
        self.password = password
        self.status  = 'connected'

    def get_info(self):
        return {
            'name'    : self.name,
            'url'     : self.url,
            'status'  : self.status,
            'username': self.username,
            'password': self.password,
        }

class OpenBISConnections(IPythonHandler):

    def post(self):
        """create a new connection

        :return: a new connection object
        """
        data = self.get_json_body()
        conn = register_connection(data)
        if conn.username and conn.password:
            try:
                conn.login()
            except Exception:
                pass
        self.get()
        return

    def get(self):
        """returns all available openBIS connections
        """

        connections= []
        for conn in openbis_connections.values():
            conn.check_status()
            connections.append(conn.get_info())

        self.write({
            'status'     : 200,
            'connections': connections,
            'cwd'        : os.getcwd()
        })
        return


class OpenBISConnectionHandler(IPythonHandler):
    """Handle the requests to /openbis/conn
    """

    def put(self, connection_name):
        """reconnect to a current connection
        :return: an updated connection object
        """
        data = self.get_json_body()

        try:
            conn = openbis_connections[connection_name]
        except KeyError:
            self.set_status(404)
            self.write({
                "reason" : 'No such connection: {}'.format(data)
            })
            return

        try:
            conn.login(data.get('username'), data.get('password'))
        except ConnectionError:
            self.set_status(500)
            self.write({
                "reason": "Could not establish connection to {}".format(connection_name)
            })
            return
        except ValueError:
            self.set_status(401)
            self.write({
                "reason": "Incorrect username or password for {}".format(connection_name)
            })
            return

        self.write({
            'status'     : 200,
            'connection': conn.get_info(),
            'cwd'        : os.getcwd()
        })

    def get(self, connection_name):
        """returns  information about a connection name
        """

        try:
            conn = openbis_connections[connection_name]
        except KeyError:
            self.set_status(404)
            self.write({
                "reason" : 'No such connection: {}'.format(data)
            })
            return

        conn.check_status()

        self.write({
            'status'     : 200,
            'connection': conn.get_info(),
            'cwd'        : os.getcwd()
        })
        return


class SampleHandler(IPythonHandler):
    """Handle the requests for /openbis/sample/connection/permId"""

    def get_datasets(self, conn, permId):
        if not conn.is_session_active():
            try:
                conn.login()
            except Exception as exc:
                self.write({
                    "reason" : 'connection to {} could not be established: {}'.format(conn.name, exc)
                })

        sample = None
        try:
            sample = conn.openbis.get_sample(permId)
        except Exception as exc:
            self.set_status(404)
            self.write({
                "reason" : 'No such sample: {}'.format(permId)
            })
        if sample is None:
            return

        datasets = sample.get_datasets().df
        datasets.replace({np.nan:None}, inplace=True)  # replace NaN with None, otherwise we cannot convert it correctly
        return datasets.to_dict(orient='records')   # is too stupid to handle NaN


    def get(self, **params):
        """Handle a request to /openbis/sample/connection_name/permId
        download the data and return a message
        """
        try:
            conn = openbis_connections[params['connection_name']]
        except KeyError:
            self.write({
                "reason" : 'connection {} was not found'.format(
                    params['connection_name']
                )
            })
            return
        
        datasets = self.get_datasets(conn, params['permId'])
        if datasets is not None:
            self.set_status(200)
            self.write({
                "dataSets": datasets
            })


class FileListHandler(IPythonHandler):

    def get(self, **params):
        """
        Returns the file list of the current working directory

        :param params:
        :return: dictionary of files, key is the fully qualified name,
                 value is the relative name (for display)
        """

        cwd = os.getcwd()
        files = {}
        for (dirpath, dirnames, filenames) in os.walk(cwd):
            if filenames:
                for filename in filenames:
                    # ignore hidden files
                    if filename.startswith('.'):
                        continue
                    # ignore hidden folders
                    if os.path.relpath(dirpath) != '.' \
                        and os.path.relpath(dirpath).startswith('.'):
                        continue
                    fqn = os.path.join(dirpath, filename)
                    files[fqn] = os.path.relpath(fqn, cwd)

        self.set_status(200)
        self.write({
            "files": files
        })


class DataSetDownloadHandler(IPythonHandler):
    """Handle the requests for /openbis/dataset/connection/permId"""


    def download_data(self, conn, permId, downloadPath=None):
        if not conn.is_session_active():
            try:
                conn.login()
            except Exception as exc:
                self.set_status(500)
                self.write({
                    "reason" : 'connection to {} could not be established: {}'.format(conn.name, exc)
                })
                return

        try:
            dataset = conn.openbis.get_dataset(permId)
        except Exception as exc:
            self.set_status(404)
            self.write({
                "reason" : 'No such dataSet found: {}'.format(permId)
            })
            return

        # dataset was found, download the data to the disk
        try: 
            destination = dataset.download(destination=downloadPath)
        except Exception as exc:
            self.set_status(500)
            self.write({
                "reason": 'Data for DataSet {} could not be downloaded: {}'.format(permId, exc)
            })
            return
            
        # return success message
        path = os.path.join(downloadPath, dataset.permId)
        self.write({
            'url'       : conn.url,
            'permId'    : dataset.permId,
            'path'      : path,
            'dataStore' : dataset.dataStore,
            'location'  : dataset.physicalData.location,
            'size'      : dataset.physicalData.size,
            'statusText': 'Data for DataSet {} was successfully downloaded to: {}.'.format(dataset.permId, path)
        })

    def get(self, **params):
        """Handle a request to /openbis/dataset/connection_name/permId
        download the data and return a message
        """

        try:
            conn = openbis_connections[params['connection_name']]
        except KeyError:
            self.set_status(404)
            self.write({
                "reason":'connection {} was not found'.format(params['connection_name'])
            })
            return
        
        results = self.download_data(conn=conn, permId=params['permId'], downloadPath=params['downloadPath'])


class DataSetTypesHandler(IPythonHandler):
    def get(self, **params):
        """Handle a request to /openbis/datasetTypes/connection_name
        """

        try:
            conn = openbis_connections[params['connection_name']]
        except KeyError:
            self.set_status(404)
            self.write({
                "reason":'connection {} was not found'.format(params['connection_name'])
            })
            return

        try:
            dataset_types = conn.openbis.get_dataset_types()
            dts = dataset_types.df.to_dict(orient='records')

            # get property assignments for every dataset-type
            # and add it to the dataset collection
            for dt in dts:
                dataset_type = conn.openbis.get_dataset_type(dt['code'])
                pa = dataset_type.get_propertyAssignments()
                pa_dict = pa.to_dict(orient='records')
                dt['propertyAssignments'] = pa_dict

            self.write({
                "dataSetTypes": dts
            })
            return

        except Exception as e:
            self.set_status(500)
            self.write({
                "reason":'Could not fetch dataset-types: {}'.format(e)
            })
            return


class DataSetUploadHandler(IPythonHandler):
    """Handle the requests for /openbis/dataset/connection"""

    def upload_data(self, conn, data):
        if not conn.is_session_active():
            try:
                conn.login()
            except Exception as exc:
                self.write({
                    "reason": 'connection to {} could not be established: {}'.format(conn.name, exc)
                })
                return

        try:
            sample = conn.openbis.get_sample(data.get('sampleIdentifier'))
        except Exception as exc:
            self.set_status(404)
            self.write({
                "reason" : 'No such sample: {}'.format(data.get('sampleIdentifier'))
            })
            return

        filenames = []
        for filename in data.get('files'):
            filename = unquote(filename)
            filenames.append(filename)

        try: 
            ds = conn.openbis.new_dataset(
                name        = data.get('name'),
                description = data.get('description'),
                type        = data.get('type'),
                sample      = sample,
                files       = filenames
            ) 
        except Exception as exc:
            self.write({
                "reason": 'Error while creating the dataset: {}'.format(exc)
            })
            return

        try:
            ds.save()
        except Exception as exc:
            self.write({
                "reason": 'Error while saving the dataset: {}'.format(exc)
            })
            return
        
            
        # return success message
        self.write({
            'status': 200,
            'statusText': 'Jupyter Notebook was successfully uploaded to: {} with permId: {}'.format(conn.name, ds.permId)
        })

    def post(self, **params):
        """Handle a request to /openbis/dataset/connection_name/permId
        download the data and return a message
        """

        try:
            conn = openbis_connections[params['connection_name']]
        except KeyError:
            self.write({
                "reason": 'connection {} was not found'.format(params['connection_name'])
            })
            return

        data = self.get_json_body()
        results = self.upload_data(conn=conn,data=data)