#!/usr/bin/env python

# Copyright (c) 2014 Anchor Systems

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

"""
NAME

    es-allocator-enable - enable an ElasticSearch cluster's automatic
                          shard allocator


SYNOPSIS

    es-allocator-enable [HOST[:PORT]]


DESCRIPTION

    Connect to the ElasticSearch HTTP server at HOST:PORT (defaults to 
    localhost:9200), and submit an API request to enable the cluster's 
    automatic shard allocator.  The automatic shard allocator is 
    responsible for transparently moving index shards between nodes in 
    order to maintain data redundancy and load balance.  You may need to 
    temporarily disable the allocator for maintenance or backups.

    You may supply the HOST:PORT of any one of the servers in your 
    ElasticSearch cluster.  The server that receives our request will 
    forward it on to its master.


RETURN VALUES

    es-allocator-enable 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


SEE ALSO

    es-allocator-disable

"""
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 [HOST[:PORT]]\n" %
                         os.path.basename(argv[0]))

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

    host = DEFAULT_HOST
    port = DEFAULT_PORT

    try:
        address = sys.argv[1]
    except IndexError:
        pass
    else:
        if address.find(":") == -1:
            host = address
        else:
            host, port = address.split(":")
            try:
                port = int(port)
            except ValueError:
                usage(argv)
                return 2

    logger.info("Enabling automatic shard allocation...")

    with Connection(host, port) as conn:
        old_status = conn.allocator_disabled()
        logger.info("Old status: %s" % old_status.upper())

        conn.put('/_cluster/settings',
                 '{"transient":{"cluster.routing.allocation.'
                                'disable_allocation":false}}')

        new_status = conn.allocator_disabled()
        logger.info("New status: %s" % new_status.upper())

        assert new_status == "enabled"

    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))
