diff --git a/django_print_sql/__init__.py b/django_print_sql/__init__.py index 37f524f..492fd63 100644 --- a/django_print_sql/__init__.py +++ b/django_print_sql/__init__.py @@ -1,3 +1,3 @@ -from .print_sql import print_sql, print_sql_decorator +from .print_sql import print_sql, print_sql_decorator, print_sql_to_file -__all__ = ['print_sql', 'print_sql_decorator'] \ No newline at end of file +__all__ = ['print_sql', 'print_sql_decorator', 'print_sql_to_file'] diff --git a/django_print_sql/print_sql.py b/django_print_sql/print_sql.py index d6f8988..4fbc82e 100644 --- a/django_print_sql/print_sql.py +++ b/django_print_sql/print_sql.py @@ -2,6 +2,11 @@ from django.db.models.sql.compiler import SQLCompiler from functools import wraps from contextlib import contextmanager +import traceback +import codecs + +import logging +logger = logging.getLogger(__name__) try: import sqlparse @@ -18,7 +23,7 @@ def format(statement, *args, **kwargs): def pprint_sql(query): statement = query[0] % query[1] - print(sqlparse.format(statement, reindent=True, keyword_case='upper')) + logger.debug(sqlparse.format(statement, reindent=True, keyword_case='upper')) original_execute_sql = SQLCompiler.execute_sql @@ -27,18 +32,21 @@ def pprint_sql(query): @contextmanager def print_sql(count_only=False): - shared_var = {'count': 0, 'total_time': 0} + shared_var = {'count': 0, 'total_time': 0, 'first_time' : 0} @wraps(original_execute_sql) def execute_sql(self, *args, **kwargs): shared_var['count'] += 1 time_begin = time.time() + if shared_var['first_time'] == 0: + shared_var['first_time'] = time_begin ret = original_execute_sql(self, *args, **kwargs) time_elapsed = time.time() - time_begin shared_var['total_time'] += time_elapsed if not count_only: pprint_sql(self.as_sql()) - print('[Time elapsed: {time:.2f}ms]\n'.format(time=time_elapsed * 1000)) + logger.debug('[{time_begin:.5f}s since start, time elapsed: {time:.2f}ms]\n\n'.format( + time_begin=time_begin-shared_var['first_time'], time=time_elapsed * 1000)) return ret # monkey patching the SQLCompiler @@ -49,7 +57,7 @@ def execute_sql(self, *args, **kwargs): # restore original execute_sql SQLCompiler.execute_sql = original_execute_sql - print('[{num_of_queries} {query_word} executed, total time elapsed {time:.2f}ms]\n'.format( + logger.debug('[{num_of_queries} {query_word} executed, total time elapsed {time:.2f}ms]\n'.format( num_of_queries=shared_var['count'], query_word='query' if shared_var['count'] == 1 else 'queries', time=shared_var['total_time'] * 1000 @@ -64,3 +72,69 @@ def wrapped(*fargs, **fkwargs): return func(*fargs, **fkwargs) return wrapped return wrapper + + +@contextmanager +def print_sql_to_file(log_file_path, trace=False, trace_base_str=None): + """ + Because the SQL and traceback log can be large and/or sensitive, + we save it to a file instead of logger.debug. + + Setting trace_base_str to the start of the absolute path + to the python source files prints the traceback info + corresponding to the point where the DB query is made. + """ + + shared_var = {'count': 0, 'total_time': 0, 'first_time' : 0} + + logf = codecs.open(log_file_path, 'w', encoding='utf-8') + + @wraps(original_execute_sql) + def execute_sql(self, *args, **kwargs): + shared_var['count'] += 1 + time_begin = time.time() + if shared_var['first_time'] == 0: + shared_var['first_time'] = time_begin + ret = original_execute_sql(self, *args, **kwargs) + time_elapsed = time.time() - time_begin + shared_var['total_time'] += time_elapsed + + if trace: + trb = traceback.extract_stack() + trb_i = len(trb) - 1 + while trb_i >=0: + if trace_base_str is not None: + # when set and the source file path starts with some string + if trb[trb_i][0].startswith(trace_base_str): + logf.write( '%s, %s, %s, %s\n' % trb[trb_i] ) + else: + logf.write( '%s, %s, %s, %s\n' % trb[trb_i] ) + trb_i -= 1 + + sql_query = self.as_sql() + sql_statement = sql_query[0] % sql_query[1] + sql_statement = (sqlparse.format(sql_statement, reindent=True, keyword_case='upper')) + logf.write(sql_statement) + logf.write("\n") + + logf.write('[{time_begin:.5f}s since start, time elapsed: {time:.2f}ms]\n\n'.format( + time_begin=time_begin-shared_var['first_time'], time=time_elapsed * 1000)) + + return ret + + # monkey patching the SQLCompiler + SQLCompiler.execute_sql = execute_sql + + yield # execute code in the `with` statement + + # restore original execute_sql + SQLCompiler.execute_sql = original_execute_sql + + logf.write('[{num_of_queries} {query_word} executed, total time elapsed {time:.2f}ms]\n\n'.format( + num_of_queries=shared_var['count'], + query_word='query' if shared_var['count'] == 1 else 'queries', + time=shared_var['total_time'] * 1000 + )) + + logf.close() +