Python custom logger - forward log events to Jaeger

jaeger_log_forwarder.py

import logging
import time
import yaml
from jaeger_client import Config
import os
import datetime

import logging
from logging import StreamHandler

class JaegerLogHandler(StreamHandler):
    # Custom StreamHandler implementation to forward python logger records to Jaeger / OpenTracing
    def __init__(self, logger, span):
        StreamHandler.__init__(self)
        self.span = span
        self.logger = logger
        # print("INIT JH")

    def emit(self, record):
        # See here https://docs.python.org/3/library/logging.html#logrecord-objects
        # print("EMIT JH")
        message = self.format(record)
        self.span.log_kv({
            "asctime": getattr(record, 'asctime', datetime.datetime.now()),
            "created": record.created,
            "filename": record.filename,
            "funcName": record.funcName,
            "levelname": record.levelname,
            "lineno": record.lineno,
            "message": message,
            "msg": record.msg
            "module": record.module,
            "msecs": record.msecs,
            "name": record.name,
            "pathname": record.pathname,
            "process": record.process,
            "processName": record.processName,
            "thread": record.thread,
            "threadName": record.threadName,
        })
    def __enter__(self):
        # print("ENTER JH BLOCK")
        self.logger.addHandler(self)

    def __exit__(self, type=None, value=None, traceback=None):
        # print("EXIT JH BLOCK")
        self.logger.removeHandler(self)
        if type:
            print(type)
            print(value)
            print(traceback)

# This is a function decorator that can also be used on class methods...
class trace_logs(object):
    def __init__(self, tracer=None, logger=None):
        """
        If there are decorator arguments, the function
        to be decorated is not passed to the constructor!
        """
        self.tracer = tracer
        self.logger = logger

    def __call__(self, func):
        """
        If there are decorator arguments, __call__() is only called
        once, as part of the decoration process! You can only give
        it a single argument, which is the function object.
        """
        def wrapped_f(*args, **kwargs):
            if not self.tracer:
                if hasattr(args[0], 'tracer'):
                    self.tracer = getattr(args[0], 'tracer')
                # else:
                #     self.tracer = get_tracer("asdf")
            with tracer.start_span(func.__name__) as span:
                if logger:
                    with JaegerLogHandler(logger=logger, span=span) as lh:
                        return func(*args, **kwargs)
                else:
                    return func(*args, **kwargs)
        return wrapped_f


# This is a python class which I want to instrument. I can instrument the individual class methods by adding the @trace_logs decorator.
# The decorator can also capture output from the python logging module.
logger = logging.getLogger('MyThing')
class MyThing:
    def __init__(self, tracer):
        self.tracer = tracer
        
    @trace_logs(logger=logger)
    def do_stuff(self, n=5):
        for i in range(n):
            logger.log(level=logging.INFO, msg="Doing the thing")
            time.sleep(5)
            
    # Use a different tracer
    # @trace_logs(logger=logger, tracer=tracer)
    # def do_stuff(self, n=5)...
    
    # If the object has self.logger and self.tracer defined then you dont need to set these in the decorator.
    # @trace_logs
    # def do_stuff(self, n=5)...


# Start the demo
if __name__ == "__main__":
    log_level = logging.DEBUG
    logger = logging.getLogger('')
    logger.handlers = []
    logging.basicConfig(format='%(asctime)s %(message)s', level=log_level)

    with open(os.path.join(os.path.dirname(__file__), '../../common/tracing_config.yaml'), 'r') as f:
        tracer_config = yaml.safe_load(f)

    config = Config(
        config=tracer_config,
        service_name='hello-jaeger',
        validate=True,
    )
    # this call also sets opentracing.tracer
    tracer = config.initialize_tracer()
    
    thing_1 = MyThing(tracer)

    while True:
        with tracer.start_span('TestSpan') as span:
            with JaegerLogHandler(span=span, logger=logger):
                logger.log(level=logging.INFO, msg='Test log (outer loop)!')

                with tracer.start_span('ChildSpan', child_of=span) as child_span:
                    with JaegerLogHandler(span=child_span, logger=logger):

                        # You can log to Jaeger using the span object
                        span.log_kv({'event': 'down below'})

                        # Or if you use the "with JaegerLogHandler..." block, all of your logger entries will also be forwarded to Jaeger
                        logger.log(level=logging.DEBUG, msg='Test debug log (inner loop)!!')
                        
                        # This function will be instrumented
                        thing_1.do_stuff(5)

                        # Regular print statements will not be forwarded...
                        print("Test stdout")

        time.sleep(2)   # yield to IOLoop to flush the spans - https://github.com/jaegertracing/jaeger-client-python/issues/50
    tracer.close()  # flush any buffered spans