OperationalLogger.py

#!/usr/bin/env python

"""Utilities to provide logging and email notification for an operational product."""

import sys
import os
import logging
import logging.handlers
import itertools
import subprocess
import socket
import atexit
import datetime
import StringIO
import inspect
import json
import traceback

import OperationalMail


def __exit_handler(logger, notify):
    """Email given addresses if there are any buffered errors."""
    for level in 'ERROR', 'CRITICAL':
        if level not in notify:
            notify[level] = []

    # add the last uncaught exception that has occurred as an error, if present
    try:
        tb_lines = traceback.format_exception(sys.last_type, sys.last_value, sys.last_traceback)
        tb_str = '\n'.join([line.strip() for line in tb_lines])
        logging.critical("Uncaught exception detected! Traceback follows:\n%s", tb_str)
    except AttributeError:
        pass

    if logger.has_errors():
        addresses = notify['ERROR']
        if logger.has_critical():
            addresses = addresses.extend(notify['CRITICAL'])
        if not addresses:
            logging.warning('There were errors, but no recipients specified. No email sent.')
        else:
            addresses = list(set(addresses))
            logger.email_log(addresses)
    elif 'DEBUG' in notify and notify['DEBUG']:
        addresses = list(set(notify['DEBUG']))
        logger.email_log(addresses, level=logging.DEBUG)
    else:
        logging.debug('No errors found, no email sent. All is well.')


class OperationalLogger:

    """Main class for Operational Logging."""

    DEFAULT_LOG_FORMAT_FILE = '[%(asctime)s][%(levelname)5s][p%(process)d][%(filename)s:%(lineno)d]  %(message)s'
    DEFAULT_LOG_FORMAT_CONSOLE = '%(asctime)s (%(levelname)5s) [%(filename)s:%(lineno)d]  %(message)s'
    DEFAULT_DATE_FORMAT_FILE = '%Y-%m-%d %H:%M:%S'
    DEFAULT_DATE_FORMAT_CONSOLE = '%H:%M:%S'
    DEFAULT_FROM_ADDRESS = 'no-reply@aer.com'

#    logger = OperationalLogger(
#        log_filename=os.path.join(log_directory, log_filename), logger_name=logger_name,
#        from_address=from_address, notify=notify, log_level=log_level, console_level=console_level,
#        log_format_file=log_format_file, date_format_file=date_format_file,
#        log_format_console=log_format_console, date_format_console=date_format_console
#    )
    def __init__(self, log_filename, logger_name, from_address, notify, log_level, console_level,
                 log_format_file, date_format_file, log_format_console, date_format_console, **kwargs):
        """Initialize logging handlers, formats, etc.

        Should not be called directly; use setup().
        """
        self.__server_name = socket.getfqdn()
        self.__debug_buffer = StringIO.StringIO()
        self.__error_buffer = StringIO.StringIO()
        self.__critical_buffer = StringIO.StringIO()
        if log_format_file is None:
            log_format_file = self.DEFAULT_LOG_FORMAT_FILE
        if log_format_console is None:
            log_format_console = self.DEFAULT_LOG_FORMAT_CONSOLE
        if date_format_file is None:
            date_format_file = self.DEFAULT_DATE_FORMAT_FILE
        if date_format_console is None:
            date_format_console = self.DEFAULT_DATE_FORMAT_CONSOLE
        if from_address is None:
            from_address = self.DEFAULT_FROM_ADDRESS
        self.__from_address = from_address

        logging.getLogger(logger_name).handlers = []
        logging.getLogger('_sub').handlers = []
        logging.basicConfig(filename=log_filename, format=log_format_file, datefmt=date_format_file, level=logging.DEBUG)
        formatter = logging.Formatter(fmt=log_format_console, datefmt=date_format_console)
        formatter_file = logging.Formatter(fmt=log_format_file, datefmt=date_format_file)

        # sub_handler is used for output from log_shell.
        # It is passed parentname and parentline in 'extra' dict,
        #   corresponding to file and line of code calling log_shell.
        formatter_sub = logging.Formatter(
            #fmt='%(asctime)s [%(parentname)s:%(parentline)d] ::  %(message)s',
            fmt='%(asctime)s         |-    %(message)s',
            datefmt='%H:%M:%S'
        )
        sub_handler = logging.handlers.WatchedFileHandler(log_filename)
        sub_handler.setLevel(logging.INFO)
        sub_handler.setFormatter(formatter_sub)
        logging.getLogger('_sub').addHandler(sub_handler)

        # Turn off propagation so root handler is not triggered.
        logging.getLogger('_sub').propagate = False

        debug_trigger = logging.StreamHandler(self.__debug_buffer)
        debug_trigger.setLevel(logging.DEBUG)
        debug_trigger.setFormatter(formatter_file)
        error_trigger = logging.StreamHandler(self.__error_buffer)
        error_trigger.setLevel(logging.ERROR)
        error_trigger.setFormatter(formatter)
        critical_trigger = logging.StreamHandler(self.__critical_buffer)
        critical_trigger.setLevel(logging.CRITICAL)
        critical_trigger.setFormatter(formatter)

        logging.getLogger(logger_name).addHandler(error_trigger)
        logging.getLogger(logger_name).addHandler(critical_trigger)
        logging.getLogger(logger_name).addHandler(debug_trigger)

        # provide a more concise format for terminal output
        if sys.stdout.isatty():
            console = logging.StreamHandler(sys.stdout)
            console.setLevel(console_level)
            console.setFormatter(logging.Formatter(
                fmt=log_format_console,
                datefmt=date_format_console,
            ))
            logging.getLogger(logger_name).addHandler(console)

            sub_handler_console = logging.StreamHandler(sys.stdout)
            sub_handler_console.setLevel(logging.INFO)
            sub_handler_console.setFormatter(formatter_sub)
            logging.getLogger('_sub').addHandler(sub_handler_console)

    def has_errors(self):
        """Return whether any ERROR-level log statements in buffer."""
        return len(self.__error_buffer.getvalue()) > 0

    def has_critical(self):
        """Return whether any CRITICAL-level log statements in buffer."""
        return len(self.__critical_buffer.getvalue()) > 0

    def email_log(self, addresses, from_address=None, level=logging.ERROR):
        """Send email to addresses with current buffered errors."""
        if from_address is None:
            from_address = self.__from_address
        logging.info("Emailing log to: %s", ', '.join(addresses))
        subject_prefix = "ERROR"
        if level == logging.CRITICAL:
            buf = self.__critical_buffer.getvalue()
            subject_prefix = "CRITICAL"
        elif level == logging.DEBUG:
            buf = self.__debug_buffer.getvalue()
            subject_prefix = "debug"
        else:
            buf = self.__error_buffer.getvalue()
        message = "<br />\n".join([self.__class__.__colorize(line) for line in buf.split('\n')])
        OperationalMail.send_status(subject='[{}] Log for {}:{}'.format(subject_prefix, self.__server_name,
                                                                           os.path.basename(sys.argv[0])),
                                    message=message, to_list=addresses, from_address=from_address,
                                    text_attachment=self.__debug_buffer.getvalue())

    @staticmethod
    def __colorize(line):
        """Add HTML markup for severity of log message."""
        if 'CRITICAL' in line:
            return '<span style="color:red">{}</span>'.format(line)
        if 'ERROR' in line:
            return '<span style="color:red">{}</span>'.format(line)
        if 'WARNING' in line:
            return '<span style="color:blue">{}</span>'.format(line)
        return line


##########################################################################################
#
# Helper functions
#
##########################################################################################


def setup(config_file=None, logger_name=None, log_filename=None, log_root_dir=None,
          from_address=None, notify=None, log_level=None, console_level=None,
          log_format_file=None, date_format_file=None, log_format_console=None, date_format_console=None):
    """Configure logger and begin logging.

    config_file: filename to load for logging config as JSON using these keywords, overwritten by args
    logger_name: name to pass to logger
    log_root_dir, log_filename: log will be created as log_root_dir/LOG_FILENAME.log
                                log_root_dir defaults to '.'
                                log_filename defaults to script name
    from_address: email 'From' address for notifications
    notify: dict containing 'LOG_LEVEL':['list', 'of', 'addresses']
    log_level: cutoff level for overall/file logging
    console_level: cutoff level for stdout
    log_format_file, date_format_file: formats to use for file logger
    log_format_console, date_format_console: formats to use for stdout logger
    """
    # Load config{} with passed args, defaulting to None.
    config = {
        'log_filename': log_filename, 'log_root_dir': log_root_dir, 'logger_name': logger_name,
        'from_address': from_address, 'notify': notify,
        'log_level': log_level, 'console_level': console_level,
        'log_format_file': log_format_file, 'date_format_file': date_format_file,
        'log_format_console': log_format_console, 'date_format_console': date_format_console,
    }
    # If config file provided, replace any None values with loaded value.
    if config_file:
        if not os.path.exists(config_file):
            logging.warn("Could not find logging config file: %s", config_file)
        else:
            logging.info("Loading logging config from: %s", config_file)
            with open(config_file, 'r') as f:
                loaded_config = json.load(f)
            for key, val in loaded_config.items():
                if key in config and config[key] is None:
                    logging.debug("- setting option from file: %s --> %s", key, val)
                    config[key] = val
    # Set reasonable defaults for remaining Nones.
    if config['log_root_dir'] is None:
        config['log_root_dir'] = os.getcwd()
#    log_directory = os.path.join(log_root_dir, 'log', datetime.datetime.utcnow().strftime('%Y-%m-%d'))
    log_directory = os.path.abspath(config['log_root_dir'])
    if not os.path.exists(log_directory):
        try:
            # default log dir permissions: SGID bit + rwxr-x---
            os.makedirs(log_directory, 02750)
        except OSError:
            raise RuntimeError('Unable to create log directory: {}'.format(log_directory))

    if config['log_filename'] is None:
        default_filename = os.path.basename(sys.argv[0]).split('.')[0] + '.log'
        config['log_filename'] = os.path.join(log_directory, default_filename)
    if config['logger_name'] is None:
        config['logger_name'] = ''
    if config['notify'] is None:
        config['notify'] = {}
    if config['log_level'] is None:
        config['log_level'] = logging.DEBUG
    if config['console_level'] is None:
        config['console_level'] = logging.INFO

#    logger = OperationalLogger(
#        log_filename=os.path.join(log_directory, log_filename), logger_name=logger_name,
#        from_address=from_address, notify=notify, log_level=log_level, console_level=console_level,
#        log_format_file=log_format_file, date_format_file=date_format_file,
#        log_format_console=log_format_console, date_format_console=date_format_console
#    )
    logger = OperationalLogger(**config)
    logging.info('Logging output from: {}'.format(' '.join(sys.argv)))
    logging.info('Log writing to: %s', config['log_filename'])
    logging.info('Process id {}, parent process id {}'.format(os.getpid(), os.getppid()))
    logging.info('-' * 60)
    #if notify:
    atexit.register(__exit_handler, logger, config['notify'])


def long_log(logging_level=logging.INFO, message=''):
    """For program output, multi-line messages, etc.

    When message has linebreaks, will break it into individual log calls for consistent formatting.
    """
    message_list = message.split("\n")
    for msg in message_list:
        logging.log(logging_level, ' ::  ' + msg)


def long_critical(message):
    """Log a multi-line CRITICAL message."""
    long_log(logging.CRITICAL, message)


def long_error(message):
    """Log a multi-line ERROR message."""
    long_log(logging.ERROR, message)


def long_warning(message):
    """Log a multi-line WARNING message."""
    long_log(logging.WARNING, message)


def long_info(message):
    """Log a multi-line INFO message."""
    long_log(logging.INFO, message)


def long_debug(message):
    """Log a multi-line DEBUG message."""
    long_log(logging.DEBUG, message)


def log_shell(arg_list, raise_err=False, return_output=False, **kwargs):
    """Pass arg_list to a subprocess call and log all output.

    raise_err: (like subprocess.check_call) if True,
               raises a RuntimeError if exit code is non-zero
    return_output: (like subprocess.check_output) returns any output
                   from shell call instead of the exit code
    """
    logging.info('Spawning child process from pid={}'.format(os.getpid()))
    logging.info('Executing: {}'.format(' '.join(arg_list)))
    try:
        my_name = os.path.basename(inspect.getfile(inspect.currentframe()))
        for f in inspect.stack():
            if my_name != os.path.basename(f[1]):
                parent_name, line_num = os.path.basename(f[1]), f[2]
                break

        proc_output = ""
        proc = subprocess.Popen(arg_list, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, **kwargs)
        for line in iter(proc.stdout.readline, b''):
            if not line:
                break
            proc_output += line
            logging.getLogger('_sub').info(line.rstrip(), extra={'parentname':parent_name, 'parentline':line_num})
        return_code = proc.wait()
        if return_code:
            if raise_err:
                logging.error('{} failed with exit code {}'.format(' '.join(arg_list), return_code))
                raise RuntimeError('{} failed with exit code {}'.format(' '.join(arg_list), return_code))
            else:
                logging.warning('{} failed with exit code {}'.format(' '.join(arg_list), return_code))
                return proc_output if return_output else return_code
    except OSError as err:
        logging.error('Error while executing %s: %s', ' '.join(arg_list), str(err))
        return err
    else:
        logging.info('Successfully finished: %s', ' '.join(arg_list))
        return proc_output if return_output else return_code


def stack_trace(log_level=logging.ERROR):
    """Log a stack trace at the given level."""
    import traceback
    tb = traceback.format_exc()
    long_log(log_level, tb)