<commit_before>"""
Check lengths of RabbitMQ queues and exit with proper return code based on
thresholds defined in provided configuration file.
"""

import logging
import os
import sys
from collections import namedtuple

import yaml
from argh import arg, dispatch_command
from pyrabbit.api import Client
from pyrabbit.http import NetworkError, HTTPError

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('check_rabbitmq_queues')
Stats = namedtuple('Stats', ['lengths', 'errors'])


DEFAULT_CONFIG = '/usr/local/etc/check_rabbitmq_queues.yml'
DEFAULT_USERNAME = 'guest'
DEFAULT_PASSWORD = '<PASSWORD>'
DEFAULT_VHOST = '/'
DEFAULT_HOSTNAME = 'localhost'
DEFAULT_PORT = 15672


def get_config(config_path):
    """
    Try to load script config from yaml file, exit with code 3 if file does not
    exists.
    :param config_path: path to config file
    :return: config dict
    """
    if not os.path.exists(config_path):
        logger.error('Configuration file %s does not exist.' % config_path)
        sys.exit(3)
    return yaml.load(open(config_path))


def get_client(cfg):
    """
    Get RabbitMQ client based on provided configuration or default settings.
    :param cfg: config dict
    :return: RabbitMQ client object
    """
    client = Client('%s:%s' % (cfg.get('host', DEFAULT_HOSTNAME),
                               cfg.get('port', DEFAULT_PORT)),
                    cfg.get('username', DEFAULT_USERNAME),
                    cfg.get('password', DEFAULT_PASSWORD))
    return client


def check_lengths(client, vhost, queues):
    """
    Check queues lengths.
    :param client: RabbitMQ client object
    :param vhost: RabbitMQ vhost name
    :param queues: queues to check
    :return: Stats(lengths=dict(queue_name: length),
                   errors=dict('critical': list queues with critical lengths,
                               'warning': list of queues with warning lengths))
    """

    stats = Stats(lengths={}, errors={'critical': [], 'warning': []})

    try:
        stdout = sys.stdout
        temp_stdout = open(os.devnull, 'w')
        sys.stdout = temp_stdout

        for queue, thresholds in queues.items():
            try:
                length = client.get_queue_depth(vhost, queue)
            except (NetworkError, HTTPError) as e:
                if isinstance(e, NetworkError):
                    warning = 'Can not communicate with RabbitMQ.'
                elif e.status == 404:
                    warning = 'Queue not found.'
                elif e.status == 401:
                    warning = 'Unauthorized.'
                else:
                    warning = 'Unhandled HTTP error, status: %s' % e.status

                stats.errors['warning'].append(queue)
                stats.lengths[queue] = warning
<commit_msg>[SYNC-1713] Change KeyError from pyrabbit to warning for nagios<commit_after>"""
Check lengths of RabbitMQ queues and exit with proper return code based on
thresholds defined in provided configuration file.
"""

import logging
import os
import sys
from collections import namedtuple

import yaml
from argh import arg, dispatch_command
from pyrabbit.api import Client
from pyrabbit.http import NetworkError, HTTPError

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('check_rabbitmq_queues')
Stats = namedtuple('Stats', ['lengths', 'errors'])


DEFAULT_CONFIG = '/usr/local/etc/check_rabbitmq_queues.yml'
DEFAULT_USERNAME = 'guest'
DEFAULT_PASSWORD = '<PASSWORD>'
DEFAULT_VHOST = '/'
DEFAULT_HOSTNAME = 'localhost'
DEFAULT_PORT = 15672


def get_config(config_path):
    """
    Try to load script config from yaml file, exit with code 3 if file does not
    exists.
    :param config_path: path to config file
    :return: config dict
    """
    if not os.path.exists(config_path):
        logger.error('Configuration file %s does not exist.' % config_path)
        sys.exit(3)
    return yaml.load(open(config_path))


def get_client(cfg):
    """
    Get RabbitMQ client based on provided configuration or default settings.
    :param cfg: config dict
    :return: RabbitMQ client object
    """
    client = Client('%s:%s' % (cfg.get('host', DEFAULT_HOSTNAME),
                               cfg.get('port', DEFAULT_PORT)),
                    cfg.get('username', DEFAULT_USERNAME),
                    cfg.get('password', <PASSWORD>_PASSWORD))
    return client


def check_lengths(client, vhost, queues):
    """
    Check queues lengths.
    :param client: RabbitMQ client object
    :param vhost: RabbitMQ vhost name
    :param queues: queues to check
    :return: Stats(lengths=dict(queue_name: length),
                   errors=dict('critical': list queues with critical lengths,
                               'warning': list of queues with warning lengths))
    """

    stats = Stats(lengths={}, errors={'critical': [], 'warning': []})

    try:
        stdout = sys.stdout
        temp_stdout = open(os.devnull, 'w')
        sys.stdout = temp_stdout

        for queue, thresholds in queues.items():
            try:
                length = client.get_queue_depth(vhost, queue)
            except (NetworkError, HTTPError, KeyError) as e:
                if isinstance(e, NetworkError):
                    warning = 'Can not communicate with RabbitMQ.'
                elif isinstance(e, KeyError):
                    warning = 'Cannot obtain queue data.'
                elif e.status == 404:
                    warning = 'Queue not found.'
                elif e.status == 401:
                    warning = 'Unauthorized.'
                else:
                    warning = 'Unhandled HTTP error, status: %s' % e.status

                stats.errors['warning'].append(queue)
                stats.lengths[queue] = warning

{'fraction_non_alphanumeric': 0.09963547995139732, 'fraction_numerical': 0.005207429265752473, 'mean_word_length': 2.468994581577363, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3196720', 'n_tokens_mistral': 1653, 'n_tokens_neox': 1531, 'n_words': 484}