sqlformatter.py

# coding: utf-8
import os
import sys
import logging
import sqlparse
from pygments import highlight
from pygments.lexers import SqlLexer
from pygments.formatters import Terminal256Formatter


class SqlFormatter(logging.Formatter):
    def __init__(self, *args, **kwargs):
        self.highlight = kwargs.pop('highlight', True)
        self.style = kwargs.pop('style', 'default')

        self.parse = kwargs.pop('parse', True)
        self.reindent = kwargs.pop('reindent', True)
        self.keyword_case = kwargs.pop('keyword_case', 'upper')

        self._lexer = SqlLexer()
        self._formatter = Terminal256Formatter(style=self.style)

        super(SqlFormatter, self).__init__(*args, **kwargs)

    def format(self, record):
        super(SqlFormatter, self).format(record)

        msg = record.sql

        if self.parse:
            msg = sqlparse.format(msg, reindent=self.reindent, keyword_case=self.keyword_case)

        if self.highlight:
            msg = highlight(msg, self._lexer, self._formatter)

        return msg


class LogDb:
    def __init__(self, name='django.db.backends', handler=None, propagate=None, **kwargs):
        self.name = name
        self.propagate = propagate

        self.handler = handler or logging.StreamHandler(sys.stdout)
        self.formatter = SqlFormatter(**kwargs)

        self.running = False

    def __call__(self, **options):
        # Toggle
        if not self.running:
            self.enable()
            self.running = True
        else:
            self.disable()
            self.running = False

        return self.running

    def enable(self):
        logger = logging.getLogger(self.name)

        self.handler.setFormatter(self.formatter)
        logger.addHandler(self.handler)
        logger.setLevel(logging.DEBUG)

        self.propagate = logger.propagate
        logger.propagate = False

    def disable(self):
        logger = logging.getLogger(self.name)

        logger.removeHandler(self.handler)
        logger.setLevel(logging.NOTSET)

        logger.propagate = self.propagate
        self.propagate = None

if os.name == 'nt':
    logdb = LogDb(highlight=False)
else:
    logdb = LogDb()