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 [email protected] -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: [email protected]
