hivefans
11/30/2016 - 9:48 AM

es-index-rate.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import argparse
import time

import requests

SLEEP_TIME = 5


def get_args():
    parser = argparse.ArgumentParser(
        description='Get the indexing and search rate for an Elasticsearch cluster.')
    parser.add_argument('--wait-time', type=int, default=5,
                        help='Time to wait to calculate rate.')
    parser.add_argument('--host', required=True,
                        help='Elasticsearch host to query.')
    parser.add_argument('--port', type=int, default=9200,
                        help='Elasticsearch port of host.')
    return parser.parse_args()


def get_node_counts(es_host):
    resp = requests.get(es_host, timeout=20)
    nodes = resp.json()
    index_ttl = 0
    search_ttl = 0
    doc_count = 0
    for node in nodes['nodes'].values():
        indices = node.get('indices', {})
        index_ttl += indices.get('indexing', {}).get('index_total', 0)
        search_ttl += indices.get('search', {}).get('query_total', 0)
        doc_count += indices.get('docs', {}).get('count', 0)
    return (index_ttl, search_ttl, doc_count)


def main():
    opts = get_args()
    es_host = 'http://{0.host}:{0.port}/_nodes/stats/indices'.format(opts)
    start = time.time()
    (index_ttl, search_ttl, doc_count) = get_node_counts(es_host)

    time.sleep(opts.wait_time)

    (index_ttl_2, search_ttl_2, doc_count_2) = get_node_counts(es_host)
    end = time.time()
    print('index rate: {:2f}'.format(
        (index_ttl_2 - index_ttl)/(end - start)))
    print('doc count rate: {:2f}'.format(
        (doc_count_2 - doc_count)/(end - start)))
    print('search rate: {:2f}'.format(
        (search_ttl_2 - search_ttl)/(end - start)))


if __name__ == '__main__':
    main()