diff --git a/README.md b/README.md index 858de827..07fafd10 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,7 @@ ex: --walreceiver, --no-walreceiver Enable/disable walreceiver checks. -w, --wrap-query Wrap query column instead of truncating. + --strip-comments Strip SQL comments from query text. --duration-mode DURATION_MODE Duration mode. Values: 1-QUERY(default), 2-TRANSACTION, 3-BACKEND. diff --git a/docs/man/pg_activity.pod b/docs/man/pg_activity.pod index b68dccbb..dda9bca2 100644 --- a/docs/man/pg_activity.pod +++ b/docs/man/pg_activity.pod @@ -272,6 +272,10 @@ required by another session. It shows following information: Wrap query column instead of truncating. +=item B<--strip-comments> + + Strip SQL comments from query text. + =item B<--min-duration=SECONDS> Don't display queries with smaller than specified duration (in seconds). diff --git a/pgactivity/cli.py b/pgactivity/cli.py index 0a782efb..000e01b4 100755 --- a/pgactivity/cli.py +++ b/pgactivity/cli.py @@ -122,6 +122,13 @@ def get_parser(prog: str | None = None) -> argparse.ArgumentParser: help="Wrap query column instead of truncating.", default=False, ) + group.add_argument( + "--strip-comments", + dest="strip_comments", + action="store_true", + help="Strip SQL comments from query text.", + default=False, + ) group.add_argument( "--duration-mode", dest="durationmode", diff --git a/pgactivity/types.py b/pgactivity/types.py index 114a13e2..40e543be 100644 --- a/pgactivity/types.py +++ b/pgactivity/types.py @@ -269,6 +269,7 @@ class UI: default=DurationMode.query, converter=DurationMode ) wrap_query: bool = False + strip_comments: bool = False sort_key: SortKey = attr.ib(default=SortKey.default(), converter=SortKey) query_mode: QueryMode = attr.ib(default=QueryMode.activities, converter=QueryMode) refresh_time: float | int = 2 diff --git a/pgactivity/ui.py b/pgactivity/ui.py index 4698b05e..26fbb61c 100644 --- a/pgactivity/ui.py +++ b/pgactivity/ui.py @@ -52,6 +52,7 @@ def main( min_duration=options.minduration, duration_mode=int(options.durationmode), wrap_query=options.wrap_query, + strip_comments=options.strip_comments, max_db_length=min(max(server_information.max_dbname_length, 8), 16), filters=data.filters, ) diff --git a/pgactivity/views.py b/pgactivity/views.py index 3f8f3846..75156b33 100644 --- a/pgactivity/views.py +++ b/pgactivity/views.py @@ -7,6 +7,7 @@ from textwrap import TextWrapper, dedent from typing import Any, Literal +import sqlparse from blessed import Terminal from . import colors, utils @@ -387,15 +388,21 @@ def get_indent(ui: UI) -> str: return " " * sum(c.min_width + 1 for c in ui.columns() if c.name != "Query") -def format_query(query: str, is_parallel_worker: bool) -> str: +def format_query( + query: str, is_parallel_worker: bool, *, strip_comments: bool = False +) -> str: r"""Return the query string formatted. >>> print(format_query("SELECT 1", True)) \_ SELECT 1 >>> format_query("SELECT 1", False) 'SELECT 1' + >>> format_query("/* prefix comment */SELECT 1 -- another comment", False, strip_comments=True) + 'SELECT 1' """ prefix = r"\_ " if is_parallel_worker else "" + if strip_comments: + query = sqlparse.format(query, strip_comments=True) return prefix + utils.clean_str(query) @@ -471,7 +478,11 @@ def cell( qwidth = width - len(indent) if qwidth > 0 and process.query is not None: - query = format_query(process.query, process.is_parallel_worker) + query = format_query( + process.query, + process.is_parallel_worker, + strip_comments=ui.strip_comments, + ) if not ui.wrap_query: query_value = query[:qwidth] diff --git a/pyproject.toml b/pyproject.toml index 38c68fb6..939c1b31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ "blessed >= 1.15.0", "humanize >= 0.5.1, !=4.12.0", "psutil >= 2.0.0", + "sqlparse >= 0.5.0", ] [project.optional-dependencies] diff --git a/tests/test_cli.py b/tests/test_cli.py index f6fcb1cd..d725591c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -15,6 +15,7 @@ def test_parser() -> None: "tempfiles": None, "walreceiver": None, "wrap_query": True, + "strip_comments": False, "durationmode": "1", "minduration": 0, "filters": [], @@ -51,3 +52,9 @@ def test_parser_flag_on() -> None: assert ns.pid is True assert ns.appname is False assert ns.wait is None + + +def test_parser_strip_comments() -> None: + parser = cli.get_parser() + ns = parser.parse_args(["--strip-comments"]) + assert ns.strip_comments is True diff --git a/tests/test_cli_help.txt b/tests/test_cli_help.txt index ea1a1815..f4c64b3f 100644 --- a/tests/test_cli_help.txt +++ b/tests/test_cli_help.txt @@ -31,6 +31,7 @@ Options: --walreceiver, --no-walreceiver Enable/disable walreceiver checks. -w, --wrap-query Wrap query column instead of truncating. + --strip-comments Strip SQL comments from query text. --duration-mode DURATION_MODE Duration mode. Values: 1-QUERY(default), 2-TRANSACTION, 3-BACKEND. diff --git a/tests/test_cli_help_py312.txt b/tests/test_cli_help_py312.txt index a7b89547..4a49943e 100644 --- a/tests/test_cli_help_py312.txt +++ b/tests/test_cli_help_py312.txt @@ -31,6 +31,7 @@ Options: --walreceiver, --no-walreceiver Enable/disable walreceiver checks. -w, --wrap-query Wrap query column instead of truncating. + --strip-comments Strip SQL comments from query text. --duration-mode DURATION_MODE Duration mode. Values: 1-QUERY(default), 2-TRANSACTION, 3-BACKEND. diff --git a/tests/test_ui.txt b/tests/test_ui.txt index 27213b35..af0963c0 100644 --- a/tests/test_ui.txt +++ b/tests/test_ui.txt @@ -47,6 +47,7 @@ Default CLI options, passed to ui.main(): ... "rds": False, ... "username": f"{postgres.info.user}", ... "wrap_query": False, +... "strip_comments": False, ... "header_show_instance": True, ... "header_show_workers": True, ... "header_show_system": True, diff --git a/tests/test_views.py b/tests/test_views.py index 4af94231..b0e3779c 100644 --- a/tests/test_views.py +++ b/tests/test_views.py @@ -63,3 +63,25 @@ def test_columns_header(capsys, term, ui, width, expected): views.columns_header(term, ui, width=width) out = capsys.readouterr()[0] assert out == expected + "\n" + + +@pytest.mark.parametrize( + "query, is_parallel_worker, strip_comments, expected", + [ + ("SELECT 1 -- trailing", False, False, "SELECT 1 -- trailing"), + ("SELECT 1 -- trailing", False, True, "SELECT 1"), + ("SELECT /* block */ 1", False, False, "SELECT /* block */ 1"), + ("SELECT 1 /* block */", False, True, "SELECT 1"), + ], +) +def test_format_query_strip_comments( + query: str, is_parallel_worker: bool, strip_comments: bool, expected: str +) -> None: + assert ( + views.format_query( + query, + is_parallel_worker, + strip_comments=strip_comments, + ) + == expected + )