update_cluster.py

#! /usr/bin/env python

import argparse
import sys
import time
import uuid
from urllib2 import HTTPError

from cloudera.director.common.client import ApiClient
from cloudera.director.v10 import (AuthenticationApi, ClustersApi, InstanceTemplatesApi)
from cloudera.director.v10.models import (Login, VirtualInstance)


def get_authenticated_client(args):
    """
    Create a new API client and authenticate against a server as admin

    @param args: dict of parsed command line arguments that
                 include server host and admin credentials
    @rtype:      ApiClient
    @return:     authenticated API client
    """

    # Start by creating a client pointing to the right server
    client = ApiClient(args.server)

    # Authenticate. This will start a session and store the cookie
    auth = AuthenticationApi(client)
    auth.login(Login(username=args.admin_username, password=args.admin_password))

    return client


def get_cluster_template(client, environment_name, deployment_name, cluster_name):
    """
    Get a current cluster template

    @param client: authenticated API client
    @param environment_name: the name of the parent environment
    @param deployment_name: the name of the parent deployment
    @param cluster_name: the name of the parent cluster
    @rtype: ClusterTemplate
    @return: current cluster template
    """
    api = ClustersApi(client)
    template = api.getTemplateRedacted(environment_name, deployment_name, cluster_name)

    return template


def update_cluster(client, environment_name, deployment_name, cluster_name, template):
    """
    Update an existing cluster

    @param client: authenticated API client
    @param environment_name: the name of the parent environment
    @param deployment_name: the name of the parent deployment
    @param cluster_name: the name of the parent cluster
    @param template: ClusterTemplate
    """
    api = ClustersApi(client)
    try:
        api.update(environment_name, deployment_name, cluster_name, template)

    except HTTPError as e:
        if e.code == 302:
            print 'Warning: a cluster with the same name already exists'
        else:
            raise e

    print "Clusters: %s" % api.list(environment_name, deployment_name)


def get_instance_template(client, environment_name, name):
    """
    Create an instance template with data from the configuration file

    @param client: authenticated API client
    @param environment_name: the name of the parent environment
    @param name: the name of the instance template
    @rtype: InstanceTemplate
    @return: instance template
    """
    api = InstanceTemplatesApi(client)
    template = api.get(environment_name, name)

    return template


def create_virtual_instance_with_random_id(client, environment_name, instance_template_name):
    """
    Create a new virtual instance object with a random ID

    @param client: authenticated API client
    @param environment_name: the name of the parent environment
    @param instance_template_name: the name of the instance template
    @rtype: VirtualInstance
    @return: virtual instance
    """
    return VirtualInstance(
            id=str(uuid.uuid4()),
            template=get_instance_template(client, environment_name, instance_template_name)
    )


def wait_for_cluster(client, environment_name, deployment_name, cluster_name):
    """
    Wait for the cluster bootstrap process to complete

    @param client: authenticated API client
    @param environment_name: the name of the parent environment
    @param deployment_name: the name of the parent deployment
    @param cluster_name: the name of the parent cluster
    """
    api = ClustersApi(client)
    stage = None
    while stage not in ['READY', 'BOOTSTRAP_FAILED']:
        sys.stdout.write(".")
        sys.stdout.flush()

        time.sleep(0.5)
        stage = api.getStatus(environment_name, deployment_name, cluster_name).stage

    print "\nCluster '%s' current stage is '%s'" % (cluster_name, stage)


def main():
    # Command line parameter
    parser = argparse.ArgumentParser(prog='update_cluster.py')
    parser.add_argument('--admin-username', default="admin",
                        help='Name of an user with administrative access (defaults to %(default)s)')
    parser.add_argument('--admin-password', default="admin",
                        help='Password for the administrative user (defaults to %(default)s)')
    parser.add_argument('--server', default="http://localhost:7189",
                        help="Cloudera Director server URL (defaults to %(default)s)")
    parser.add_argument('--environment', help="the name of the parent environment")
    parser.add_argument('--deployment', help="the name of the parent deployment")
    parser.add_argument('--cluster', help="the name of the parent cluster")
    parser.add_argument('--instance-template', help="the name of the instance template")
    parser.add_argument('--cluster-size', help="cluster size", default=3)
    args = parser.parse_args()

    # Log in to the API
    client = get_authenticated_client(args)

    # Get current cluster template
    template = get_cluster_template(client, args.environment, args.deployment, args.cluster)
    # print(vars(template))

    # Edit cluster template
    cluster_size = int(args.cluster_size)
    current_cluster_size = len(template.virtualInstanceGroups['workers'].virtualInstances)
    if cluster_size > current_cluster_size:
        # print 'Add cluster'
        num = cluster_size - current_cluster_size
        for i in range(num):
            instance = create_virtual_instance_with_random_id(client, args.environment, args.instance_template)
            template.virtualInstanceGroups['workers'].virtualInstances.append(instance)
    elif (cluster_size < current_cluster_size) and (cluster_size >= 3):
        # print 'Remove cluster'
        num = current_cluster_size - cluster_size
        for i in range(num):
            template.virtualInstanceGroups['workers'].virtualInstances.pop()
    else:
        print 'nop'
        return 0

    # print(vars(template.virtualInstanceGroups['workers']))

    print 'Update a cluster ...'
    update_cluster(client, args.environment, args.deployment, args.cluster, template)

    # print 'Waiting for the cluster to be ready. Check the web interface for details.'
    wait_for_cluster(client, args.environment, args.deployment, args.cluster)

    return 0


if __name__ == '__main__':
    try:
        sys.exit(main())

    except HTTPError as e:
        print e.read()
        raise e