Merge pull request #25 from dataiku/feature/multi-cluster

Feature/multi cluster

Author: cstenac

dataikuapi/dss/admin.py

@@ -567,3 +567,96 @@ def set_definition(self, definition):
         return self.client._perform_empty(
             "PUT", "/admin/globalAPIKeys/%s" % self.key,
             body = definition)
+
+class DSSCluster(object):
+    """
+    A cluster on the DSS instance
+    """
+    def __init__(self, client, cluster_id):
+        self.client = client
+        self.cluster_id = cluster_id
+    
+    ########################################################
+    # Cluster deletion
+    ########################################################
+    
+    def delete(self):
+        """
+        Delete the cluster (not stopping it)
+        """
+        self.client._perform_empty(
+            "DELETE", "/admin/clusters/%s" % (self.cluster_id))
+
+        
+    ########################################################
+    # Cluster description
+    ########################################################
+    
+    def get_definition(self):
+        """
+        Get the cluster's definition
+        
+        Returns:
+            the cluster definition, as a JSON object
+        """
+        return self.client._perform_json(
+            "GET", "/admin/clusters/%s" % (self.cluster_id))
+
+    def set_definition(self, cluster):
+        """
+        Set the cluster's definition. The definition should come from a call to the get_definition()
+        method. 
+
+        Fields that can be updated :
+
+        * cluster.permissions, cluster.usableByAll, cluster.owner
+        * cluster.params
+        
+        :param cluster: a cluster definition
+
+        Returns:
+            the updated cluster definition, as a JSON object
+        """
+        return self.client._perform_json(
+            "PUT", "/admin/clusters/%s" % (self.cluster_id), body=cluster)
+
+    def get_status(self):
+        """
+        Get the cluster's status and usages
+        
+        Returns:
+            the cluster status, as a JSON object
+        """
+        return self.client._perform_json(
+            "GET", "/admin/clusters/%s/status" % (self.cluster_id))
+
+   
+    ########################################################
+    # Cluster actions
+    ########################################################
+
+    def start(self):
+        """
+        Starts the cluster
+        """
+        resp = self.client._perform_json(
+            "POST", "/admin/clusters/%s/start" % (self.cluster_id))
+        if resp is None:
+            raise Exception('Env update returned no data')
+        if resp.get('messages', {}).get('error', False):
+            raise Exception('Cluster operation failed : %s' % (json.dumps(resp.get('messages', {}).get('messages', {}))))
+        return resp
+
+
+    def stop(self):
+        """
+        Stops the cluster
+        """
+        resp = self.client._perform_json(
+            "POST", "/admin/clusters/%s/stop" % (self.cluster_id))
+        if resp is None:
+            raise Exception('Env update returned no data')
+        if resp.get('messages', {}).get('error', False):
+            raise Exception('Cluster operation failed : %s' % (json.dumps(resp.get('messages', {}).get('messages', {}))))
+        return resp
+

dataikuapi/dssclient.py

@@ -6,7 +6,7 @@
 from .dss.future import DSSFuture
 from .dss.project import DSSProject
 from .dss.plugin import DSSPlugin
-from .dss.admin import DSSUser, DSSGroup, DSSConnection, DSSGeneralSettings, DSSCodeEnv, DSSGlobalApiKey
+from .dss.admin import DSSUser, DSSGroup, DSSConnection, DSSGeneralSettings, DSSCodeEnv, DSSGlobalApiKey, DSSCluster
 from .dss.meaning import DSSMeaning
 from .dss.sqlquery import DSSSQLQuery
 from .dss.notebook import DSSNotebook
@@ -394,6 +394,55 @@ def create_code_env(self, env_lang, env_name, deployment_mode, params=None):
             raise Exception('Env creation failed : %s' % (json.dumps(resp.get('messages', {}).get('messages', {}))))
         return DSSCodeEnv(self, env_lang, env_name)
 
+    ########################################################
+    # Clusters
+    ########################################################
+
+    def list_clusters(self):
+        """
+        List all clusters setup on the DSS instance
+
+        Returns:
+            List clusters (name, type, state)
+        """
+        return self._perform_json(
+            "GET", "/admin/clusters/")
+
+    def get_cluster(self, cluster_id):
+        """
+        Get a handle to interact with a specific cluster
+        
+        Args:
+            name: the name of the desired cluster
+        
+        Returns:
+            A :class:`dataikuapi.dss.admin.DSSCluster` cluster handle
+        """
+        return DSSCluster(self, cluster_id)
+
+    def create_cluster(self, cluster_name, cluster_type='manual', params=None):
+        """
+        Create a cluster, and return a handle to interact with it
+
+        :param cluster_name: the name of the new cluster
+        :param cluster_type: the type of the new cluster
+        :param params: the parameters of the new cluster, as a JSON object
+        
+        :returns: A :class:`dataikuapi.dss.admin.DSSCluster` cluster handle
+        
+        """
+        definition = {}
+        definition['name'] = cluster_name
+        definition['type'] = cluster_type
+        definition['params'] = params if params is not None else {}
+        resp = self._perform_json(
+               "POST", "/admin/clusters/", body=definition)
+        if resp is None:
+            raise Exception('Cluster creation returned no data')
+        if resp.get('messages', {}).get('error', False):
+            raise Exception('Cluster creation failed : %s' % (json.dumps(resp.get('messages', {}).get('messages', {}))))
+        return DSSCluster(self, resp['id'])
+
     ########################################################
     # Global API Keys
     ########################################################