nagios-check-solr.py

#!/usr/bin/env python

import optparse

import urllib
import urllib2
import urlparse

import xml.etree.ElementTree as ET


NAGIOS_OK = 0
NAGIOS_WARNING = 1
NAGIOS_CRITICAL = 2
NAGIOS_UNKNOWN = 3

def nagios_name(n):
    return ['OK', 'WARNING', 'CRITICAL', 'UNKNOWN'][n]


def main():
    checks = [
        SelectCheck('DEFAULT', rows=0),
        SelectCheck('bibcodes', qt='bibcodes', rows=0),
        SelectCheck('proposal-codes', qt='proposal-codes', rows=0)]

    parser = optparse.OptionParser(usage='%prog OPTIONS SOLR',
            epilog='Check names are: %s' % (', '.join( check.name for check in checks ),))
    parser.add_option('-w', '--warning', action='append', dest='warning_thresholds')
    parser.add_option('-c', '--critical', action='append', dest='critical_thresholds')

    options, args = parser.parse_args()

    try:
        solr_address, = args

    except ValueError:
        parser.error('Specify SOLR index address')

    try:
        warning_thresholds = dict( threshold_option(s)
                for s in options.warning_thresholds or [] )
        critical_thresholds = dict( threshold_option(s)
                for s in options.critical_thresholds or [] )

    except ValueError:
        parser.error('Threshold format is CHECK_NAME=N')

    results = [
            SelectCheckResult(check,
                warning_thresholds.pop(check.name, None),
                critical_thresholds.pop(check.name, None))
            for check in checks ]

    for name in list(warning_thresholds)+list(critical_thresholds):
        parser.error('unknown check name %r' % (name,))

    for result in results:
        result(solr_address)

    max_status = max( result.status for result in results )
    summary = '/'.join( result.data for result in results )

    print nagios_name(max_status), summary
    for result in results:
        print result.desc
    print '|'
    for result in results:
        print '%s=%s' % (result.name, result.data)

    exit(max_status)


class SelectCheck(object):
    def __init__(self, name, *a, **kw):
        self.name = name
        self.params = kw.items()
        self.params.extend(a)


class SelectCheckResult(object):
    def __init__(self, check, warning, critical):
        self.check = check
        self.warning = warning
        self.critical = critical

    status = None
    error = None
    hits = None

    @property
    def name(self):
        return self.check.name

    def __call__(self, solr_address):
        try:
            response = get_select(solr_address, *self.check.params)

        except IOError, e:
            self.status = NAGIOS_UNKNOWN
            self.error = str(e)
            if hasattr(e, 'url'):
                self.error += ' (%s)' % (e.url,)

        else:
            self.hits = get_select_hits(response)

            if self.critical and self.hits < self.critical:
                self.status = NAGIOS_CRITICAL

            elif self.warning and self.hits < self.warning:
                self.status = NAGIOS_WARNING

            else:
                self.status = NAGIOS_OK

    @property
    def desc(self):
        if self.status == NAGIOS_UNKNOWN:
            return '%s ERROR %s' % (self.name, self.error)
        if self.status == NAGIOS_CRITICAL:
            return '%s %s %s < %s' % (self.name, nagios_name(self.status),
                    self.hits, self.critical)
        if self.status == NAGIOS_WARNING:
            return '%s %s %s < %s' % (self.name, nagios_name(self.status),
                    self.hits, self.warning)
        if self.status == NAGIOS_OK:
            return '%s %s %s' % (self.name, nagios_name(self.status),
                    self.hits)

    @property
    def data(self):
        if self.hits is None:
            return '?'
        return str(self.hits)


def get_admin_stats(solr_address):
    url = urlparse.urljoin(solr_address, 'admin/stats.jsp')
    fh = urllib2.urlopen(url)
    return ET.parse(fh)

def get_select(solr_address, *a, **kw):
    params = kw.items()
    params.extend(a)
    url = urlparse.urljoin(solr_address, 'select')
    url += '?'+urllib.urlencode(params)
    fh = urllib2.urlopen(url)
    return ET.parse(fh)

def get_select_hits(response):
    for el in response.findall('./result'):
        if el.get('name') == 'response':
            return int(el.get('numFound'))

def threshold_option(s):
    if '=' in s:
        name, count = s.split('=', 1)
    else:
        name, count = 'DEFAULT', s
    return name, int(count)


if __name__ == '__main__':
    main()