AIL-framework/bin/pubsublogger/subscriber.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
:mod:`subscriber` -- Subscribe to a redis channel and gather logging messages.

To use this module, you have to define at least a channel name.
"""


import redis
from logbook import Logger
import ConfigParser
from logbook import NestedSetup
from logbook import NullHandler
from logbook import TimedRotatingFileHandler
from logbook import MailHandler
import os

# use a TCP Socket by default
use_tcp_socket = True

# default config for a UNIX socket
unix_socket = '/tmp/redis.sock'
# default config for a TCP socket
hostname = 'localhost'
port = 6379

pubsub = None
channel = None

# Required only if you want to send emails
dest_mails = []
smtp_server = None
smtp_port = 0
src_server = None


def setup(name, path='log', enable_debug=False):
    """
    Prepare a NestedSetup.

    :param name: the channel name
    :param path: the path where the logs will be written
    :param enable_debug: do we want to save the message at the DEBUG level

    :return a nested Setup
    """
    path_tmpl = os.path.join(path, '{name}_{level}.log')
    info = path_tmpl.format(name=name, level='info')
    warn = path_tmpl.format(name=name, level='warn')
    err = path_tmpl.format(name=name, level='err')
    crit = path_tmpl.format(name=name, level='crit')
    # a nested handler setup can be used to configure more complex setups
    setup = [
        # make sure we never bubble up to the stderr handler
        # if we run out of setup handling
        NullHandler(),
        # then write messages that are at least info to to a logfile
        TimedRotatingFileHandler(info, level='INFO', encoding='utf-8',
                                 date_format='%Y-%m-%d'),
        # then write messages that are at least warnings to to a logfile
        TimedRotatingFileHandler(warn, level='WARNING', encoding='utf-8',
                                 date_format='%Y-%m-%d'),
        # then write messages that are at least errors to to a logfile
        TimedRotatingFileHandler(err, level='ERROR', encoding='utf-8',
                                 date_format='%Y-%m-%d'),
        # then write messages that are at least critical errors to to a logfile
        TimedRotatingFileHandler(crit, level='CRITICAL', encoding='utf-8',
                                 date_format='%Y-%m-%d'),
    ]
    if enable_debug:
        debug = path_tmpl.format(name=name, level='debug')
        setup.insert(1, TimedRotatingFileHandler(debug, level='DEBUG',
                     encoding='utf-8', date_format='%Y-%m-%d'))
    if src_server is not None and smtp_server is not None \
            and smtp_port != 0 and len(dest_mails) != 0:
        mail_tmpl = '{name}_error@{src}'
        from_mail = mail_tmpl.format(name=name, src=src_server)
        subject = 'Error in {}'.format(name)
        # errors should then be delivered by mail and also be kept
        # in the application log, so we let them bubble up.
        setup.append(MailHandler(from_mail, dest_mails, subject,
                                 level='ERROR', bubble=True,
                                 server_addr=(smtp_server, smtp_port)))

    return NestedSetup(setup)


def mail_setup(path):
    """
    Set the variables to be able to send emails.

    :param path: path to the config file
    """
    global dest_mails
    global smtp_server
    global smtp_port
    global src_server
    config = ConfigParser.RawConfigParser()
    config.readfp(path)
    dest_mails = config.get('mail', 'dest_mail').split(',')
    smtp_server = config.get('mail', 'smtp_server')
    smtp_port = config.get('mail', 'smtp_port')
    src_server = config.get('mail', 'src_server')


def run(log_name, path, debug=False, mail=None):
    """
    Run a subscriber and pass the messages to the logbook setup.
    Stays alive as long as the pubsub instance listen to something.

    :param log_name: the channel to listen to
    :param path: the path where the log files will be written
    :param debug: True if you want to save the debug messages too
    :param mail: Path to the config file for the mails

    """
    global pubsub
    global channel
    channel = log_name
    if use_tcp_socket:
        r = redis.StrictRedis(host=hostname, port=port)
    else:
        r = redis.StrictRedis(unix_socket_path=unix_socket)
    pubsub = r.pubsub()
    pubsub.psubscribe(channel + '.*')

    logger = Logger(channel)
    if mail is not None:
        mail_setup(mail)
    if os.path.exists(path) and not os.path.isdir(path):
        raise Exception("The path you want to use to save the file is invalid (not a directory).")
    if not os.path.exists(path):
        os.mkdir(path)
    with setup(channel, path, debug):
        for msg in pubsub.listen():
            if msg['type'] == 'pmessage':
                level = msg['channel'].split('.')[1]
                message = msg['data']
                try:
                    message = message.decode('utf-8')
                except:
                    pass
                logger.log(level, message)


def stop():
    """
    Unsubscribe to the channel, stop the script.
    """
    pubsub.punsubscribe(channel + '.*')