Tuesday, December 1, 2015

PyMongo query logging

I needed to see what queries are sent to the server in my MongoEngine application.

I made one logger here: https://github.com/warvariuc/python-mongo-logger based on https://gist.github.com/kesor/1589672

Recently I improved the solution, but did not commit to the repo, because I put this code directly into the project instead of having git+git://github.com/warvariuc/python-mongo-logger
 in requirements.pip file:

from __future__ import absolute_import, print_function, unicode_literals, division

import logging
import time
import traceback
import inspect

from pymongo.mongo_client import MongoClient
from bson import json_util


logger = logging.getLogger('mongologger')


def activate(until_modules=('pymongo', 'mongoengine'), stack_size=3):
    """Activate Mongo-Logger.
    Args:
        until_modules (list): list of top level module names until which the stack should be shown;
          pass an empty sequence to show the whole stack
        stack_size (int): how many frames before any of `modules` was entered to show; pass
          -1 to show the whole stack or 0 to show no stack
    """
    # monkey-patch methods to record messages
    MongoClient._send_message_with_response = _instrument(
        MongoClient._send_message_with_response, until_modules, stack_size)
    return logger


def _instrument(original_method, until_modules, stack_size):
    """Monkey-patch the given pymongo function which sends queries to MongoDB.
    """
    def _send_message_with_response(client, operation, read_preference=None,
                                    exhaust=False, address=None):
        start_time = time.time()
        result = original_method(client, operation, read_preference=None,
                                    exhaust=False, address=None)
        duration = time.time() - start_time
        try:
            stack = ('\n' + ''.join(get_stack(until_modules, stack_size))).rstrip()
            logger.info('%.3f %s %s %s%s', duration, operation.name, operation.ns,
                        json_util.dumps(operation.spec, ensure_ascii=False), stack)
        except Exception as exc:
            logger.info('%.3f *** Failed to log the query *** %s', duration, exc)
        return result

    return _send_message_with_response


def get_stack(until_modules, stack_size):
    """
    """
    if not stack_size:
        return []
    frames = inspect.stack()[2:]
    frame_index = None
    for i, (frame, _, _, _, _, _) in enumerate(frames):
        module_name, _, _ = frame.f_globals['__name__'].partition('.')
        if module_name in until_modules:
            frame_index = i
        elif frame_index is not None:
            # found first frame before the needed module frame was entered
            break

    if frame_index is not None:
        del frames[:frame_index + 1]
        if stack_size >= 0:
            del frames[stack_size:]

    stack = [(filename, lineno, name, lines[0])
             for frame, filename, lineno, name, lines, _ in frames]

    return traceback.format_list(stack)

You activate the logger in the logging configuration like this:

import mongologger
LOGGING['loggers'][mongologger.activate(stack_size=0).name] = {
    'level': 'INFO',
    'handlers': ['stdout'],
    'propagate': False,
}


No comments:

Post a Comment