# -*- coding: utf-8 -*-
##
##
## This file is part of Indico.
## Copyright (C) 2002 - 2014 European Organization for Nuclear Research (CERN).
##
## Indico is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 3 of the
## License, or (at your option) any later version.
##
## Indico is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with Indico; if not, see <http://www.gnu.org/licenses/>.

from __future__ import absolute_import

import logging
import os
import pprint
import time

from sqlalchemy.engine import Engine
from sqlalchemy.event import listens_for

# Optional dependencies for pretty SQL logging
try:
    import sqlparse
except ImportError:
    sqlparse = None

try:
    from pygments import highlight
    from pygments.lexers import SqlLexer, PythonLexer
    from pygments.formatters import Terminal256Formatter
except ImportError:
    has_pygments = False
else:
    has_pygments = True


def _prettify_sql(statement):
    if sqlparse:
        statement = sqlparse.format(statement, keyword_case='upper', reindent=True)
    statement = '    ' + statement.replace('\n', '\n    ')
    if not has_pygments or os.environ.get('INDICO_COLORED_LOG') != '1':
        return statement
    return highlight(statement, SqlLexer(), Terminal256Formatter(style='native'))


def _prettify_params(args):
    args = pprint.pformat(args)
    args = '    ' + args.replace('\n', '\n    ')
    if not has_pygments or os.environ.get('INDICO_COLORED_LOG') != '1':
        return args
    return highlight(args, PythonLexer(), Terminal256Formatter(style='native')).rstrip()


def apply_db_loggers(debug=False):
    if not debug:
        return
    from indico.core.logger import Logger

    logger = Logger.get('db')
    logger.setLevel(logging.DEBUG)

    @listens_for(Engine, 'before_cursor_execute')
    def before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
        context._query_start_time = time.time()
        logger.debug('Start Query:\n{}\n{}'.format(_prettify_sql(statement),
                                                   _prettify_params(parameters) if parameters else '').rstrip())

    @listens_for(Engine, 'after_cursor_execute')
    def after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
        total = time.time() - context._query_start_time
        logger.debug('Query complete; total time: {}'.format(total))

@@ -856,16 +856,33 @@
 ort time
+%0Aimport traceback
 %0A%0Afrom s
@@ -1961,184 +1961,1187 @@
 ()%0A%0A
-%0Adef apply_db_loggers(debug=False):%0A    if not debug:%0A        return%0A    from indico.core.logger import Logger%0A%0A    logger = Logger.get('db')%0A    logger.setLevel(logging.DEBUG)
+def _prettify_traceback(args):%0A    args = pprint.pformat(args)%0A    args = args.replace('%5Cn', '')%0A    if not has_pygments or os.environ.get('INDICO_COLORED_LOG') != '1':%0A        return args%0A    return highlight(args, PythonLexer(), Terminal256Formatter(style='native')).strip()%0A%0A%0Adef apply_db_loggers(debug=False):%0A    if not debug:%0A        return%0A    from indico.core.logger import Logger%0A%0A    logger = Logger.get('db')%0A    logger.setLevel(logging.DEBUG)%0A%0A    def print_sql_line():%0A        import indico%0A%0A        indico_path = os.path.dirname(os.path.abspath(indico.__file__))%0A        root_path = %22%7B%7D/%22.format(os.path.dirname(indico_path))%0A        stack = traceback.extract_stack()%0A        for item in reversed(stack):%0A            if item%5B0%5D.startswith(indico_path) and 'logging' not in item%5B0%5D and 'sqlalchemy' not in item%5B0%5D:%0A                module_name = os.path.splitext(item%5B0%5D.replace(root_path, ''))%5B0%5D.replace(os.sep, '.')%0A                logger.debug('%5Cn%7B%7D:%7B%7D %7B%7D'.format(_prettify_traceback(module_name),%0A                                         _prettify_traceback(item%5B1%5D),%0A                                         _prettify_traceback(item%5B2%5D)))%0A                break%0A
 %0A%0A  
@@ -3270,32 +3270,57 @@
 , executemany):%0A
+        print_sql_line()%0A
         context.

Print SQL line code in log