1 files changed, 19 insertions, 46 deletions
diff --git a/debug_toolbar/panels/sql.py b/debug_toolbar/panels/sql.py
index 0c9bc61..e1a2b7d 100644
--- a/debug_toolbar/panels/sql.py
+++ b/debug_toolbar/panels/sql.py
@@ -13,6 +13,7 @@ from django.utils.encoding import force_unicode
 from django.utils.hashcompat import sha_constructor
 
 from debug_toolbar.panels import DebugPanel
+from debug_toolbar.utils import sqlparse
 
 # Figure out some paths
 django_path = os.path.realpath(os.path.dirname(django.__file__))
@@ -20,48 +21,8 @@ socketserver_path = os.path.realpath(os.path.dirname(SocketServer.__file__))
 
 # TODO:This should be set in the toolbar loader as a default and panels should
 # get a copy of the toolbar object with access to its config dictionary
-SQL_WARNING_THRESHOLD = getattr(settings, 'DEBUG_TOOLBAR_CONFIG', {}).get('SQL_WARNING_THRESHOLD', 500)
-
-# Note: This isn't intended to catch ALL possible SQL keywords, just a good common set.
-# Note: Subsets are listed last to avoid matching a subset of a keyword.  This
-# whole thing could be greatly improved but for now this works.
-SQL_KEYWORDS = (
-    'ALTER',
-    'AND',
-    'ASC',
-    'AS',
-    'AVG',
-    'COUNT',
-    'CREATE',
-    'DESC',
-    'DELETE',
-    'DISTINCT',
-    'DROP',
-    'FROM',
-    'GROUP BY',
-    'HAVING',
-    'INNER JOIN',
-    'INSERT',
-    'IN',
-    'LEFT OUTER JOIN',
-    'LIKE',
-    'LIMIT',
-    'MAX',
-    'MIN',
-    'OFFSET',
-    'ON',
-    'ORDER BY',
-    'OR',
-    'SELECT',
-    'SET',
-    'STDDEV_POP',
-    'STDDEV_SAMP',
-    'SUM',
-    'UPDATE',
-    'VAR_POP',
-    'VAR_SAMP',
-    'WHERE',
-)
+SQL_WARNING_THRESHOLD = getattr(settings, 'DEBUG_TOOLBAR_CONFIG', {}) \
+                            .get('SQL_WARNING_THRESHOLD', 500)
 
 def tidy_stacktrace(strace):
     """
@@ -170,8 +131,20 @@ def ms_from_timedelta(td):
     """
     return (td.seconds * 1000) + (td.microseconds / 1000.0)
 
-def reformat_sql(sql):
-    for kwd in SQL_KEYWORDS:
-        sql = sql.replace(kwd, '<strong>%s</strong>' % (kwd,))
-    return sql
+class BoldKeywordFilter(sqlparse.filters.Filter):
+    """sqlparse filter to bold SQL keywords"""
+    def process(self, stack, stream):
+        """Process the token stream"""
+        for token_type, value in stream:
+            is_keyword = token_type in sqlparse.tokens.Keyword
+            if is_keyword:
+                yield sqlparse.tokens.Text, '<strong>'
+            yield token_type, value
+            if is_keyword:
+                yield sqlparse.tokens.Text, '</strong>'
 
+def reformat_sql(sql):
+    stack = sqlparse.engine.FilterStack()
+    stack.preprocess.append(BoldKeywordFilter()) # add our custom filter
+    stack.postprocess.append(sqlparse.filters.SerializerUnicode()) # tokens -> strings
+    return ''.join(stack.run(sql))