[GitHub] [airflow] kaxil commented on a change in pull request #8754: Add SQL query tracking for pytest

[GitBox <gi...@apache.org>]

kaxil commented on a change in pull request #8754:
URL: https://github.com/apache/airflow/pull/8754#discussion_r421442805



##########
File path: scripts/perf/perf_kit/sqlalchemy.py
##########
@@ -15,31 +15,59 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import contextlib
 import os
 import time
 import traceback
+from typing import Callable
 
 from sqlalchemy import event
 
 
-@contextlib.contextmanager
-def trace_queries(display_time=True, display_trace=True, display_sql=False, display_parameters=True):
+def _pretty_format_sql(text: str):
+    import pygments
+
+    from pygments.formatters.terminal import TerminalFormatter
+    from pygments.lexers.sql import SqlLexer
+    text = pygments.highlight(
+        code=text, formatter=TerminalFormatter(), lexer=SqlLexer()
+    ).rstrip()
+    return text
+
+
+class TraceQueries:
     """
-    Tracking SQL queries in a code block. The result is displayed directly on the screen - ``print``
+    Tracking SQL queries in a code block.
 
+    :param display_no: If True, displays the query number.
     :param display_time: If True, displays the query execution time.
     :param display_trace: If True, displays the simplified (one-line) stack trace
     :param display_sql: If True, displays the SQL statements
     :param display_parameters: If True, display SQL statement parameters
-    :return:
+    :param print_fn: The function used to display the text. By default,``builtins.print``
     """
-    import airflow.settings
-
-    def before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
+    def __init__(
+        self,
+        *,

Review comment:
       what's this for?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org