#!/usr/bin/env python

# Copyright (c) 2014 Anchor Systems

# https://github.com/anchor/elasticsearch-scripts

"""
NAME

    es-concurrent-rebalance - display or set the number of shards
                              Elasticsearch is allowed to balance
                              at one time.

SYNOPSIS

    es-concurrent-rebalance [HOST[:PORT]]
    
    es-concurrent-rebalance maximum_shards [HOST[:PORT]]


DESCRIPTION

    Run with no arguments to display the current value of this setting.
    Run with one integral argument to set it. Most of the time it should
    probably be set to 0. 

RETURN VALUES

    es-concurrent-rebalance will return zero on success; non-zero on 
    failure.  Success is defined by the underlying ElasticSearch API.  
    See BACKGROUND INFORMATION.


BACKGROUND INFORMATION

    - http://www.elasticsearch.org/guide/reference/modules/cluster.html
    - http://www.elasticsearch.org/guide/reference/index-modules/allocation.html

"""
DEFAULT_HOST = "localhost"
DEFAULT_PORT = 9200

from esadmin import Connection
import logging
import os
import socket
import sys

logger = logging.getLogger(__name__)

def usage(argv):
    print >>sys.stderr, ("Usage:\n\t%s [ maximum_shards ] \n" %
                         os.path.basename(argv[0]))

def main(argv=None):
    if argv is None:
        argv = sys.argv

    host = DEFAULT_HOST
    port = DEFAULT_PORT

    new_max = None

    try:
        new_max = int(sys.argv[1])
    except ValueError:
        usage()
        sys.exit(2)
    except IndexError:
        pass

    with Connection(host, port) as conn:
        setting_key = "cluster.routing.allocation.cluster_concurrent_rebalance"
        settings = conn.get('/_cluster/settings')
        old_max = None
        for group in ("persistent", "transient"):
            try:
                old_max = settings[group][setting_key]
            except KeyError:
                pass

        old_max_report = "%s" % (
                str(old_max) if 
                    old_max is not None else
                "Could not retrieve %s via the settings API." % setting_key
        )

        if new_max is None:
            print(old_max_report)
            sys.exit(0)

        logger.info("Old maximium concurrent rebalance: %s" % old_max_report)

        conn.put('/_cluster/settings',
                 '{"transient":{"%s":%d}}' % (
                    setting_key,
                    new_max,
                )
        )

        logger.info("New maximum concurrent rebalance: %d" % new_max)

    logger.info("Done")
    return 0

if __name__ == '__main__':
    hostname = socket.gethostname()
    fmt = ('[%s] [%%(process)d]: [%%(levelname)s] '
           '%%(message)s' % hostname)
    logging.basicConfig(level=logging.INFO, format=fmt)

    sys.exit(main(sys.argv))
