escount.py

#!/usr/local/bin/python3
import urllib.request
import json
import sys
import os

URL = 'http://0.0.0.0:9200/_stats'
# COUNT_FILE = os.environ['HOME'] + '/tmp/es_count.tmp'
# COUNT_OUTPUT = os.environ['HOME'] + '/tmp/es_output.tmp'
COUNT_FILE = '/Users/cmoultrie/tmp/es_count.tmp'
COUNT_OUTPUT = '/Users/cmoultrie/tmp/es_output.tmp'


def get_count():
    """
    Attempts to get the current count from ES, bailing if no bueno
    """
    try:
        # Timeout of 1s, sometimes ES gets hosed and we block the entire Screen
        # Session
        stats_conn = urllib.request.urlopen(URL)
        stats = json.loads(stats_conn.read().decode('utf-8'))
        # All docs everywhere
        # count = stats['primaries']['docs']['count']
        # Just the pindrop index
        count = stats['indices']['pindrop']['primaries']['docs']['count']

        count_format = "{:,}".format(count)
    except Exception:
        return None

    return [count, count_format]


def get_diff(current):
    """
    Checks for increase/decrease in Document Count and formats the statusline
    appropriately
    """
    color = 'G'
    diff_fmt = ''

    previous = None
    if os.path.exists(COUNT_FILE):
        with open(COUNT_FILE, 'r') as counted:
            previous = int(counted.read())

    if previous is not None:
        if previous < current:
            color = 'G'
            diff_fmt = "+{:,}".format(current - previous)
        elif previous > current:
            color = 'R'
            diff_fmt = "-{:,}".format(previous - current)

        if diff_fmt != '':
            diff_fmt = " \005{c}%s\005{g}" % diff_fmt

    if previous != current:
        with open(COUNT_FILE, 'w+') as counted:
            counted.write("%d" % current)

    return [color, diff_fmt]

if __name__ == "__main__":
    counts = get_count()
    if counts is None:
        with open(COUNT_OUTPUT, 'w+') as output:
            output.write('')
        sys.exit(0)

    (count_num, count_fmt) = counts
    (text_color, diff) = get_diff(count_num)

    fmt_vars = (text_color, count_fmt, diff)

    output_text = "[ \005{%s}%s\005{g}%s ]" % fmt_vars
    with open(COUNT_OUTPUT, 'w+') as output:
        output.write(output_text)
    sys.stdout.write(output_text)