diff --git a/AUTHORS b/AUTHORS index bf9a3b226..aef249cef 100644 --- a/AUTHORS +++ b/AUTHORS @@ -147,6 +147,7 @@ Contributors: * Devadathan M B (devadathanmb) * Charalampos Stratakis * Laszlo Bimba (bimlas) + * Anjanna Creator: -------- diff --git a/changelog.rst b/changelog.rst index 78e8ee394..020113660 100644 --- a/changelog.rst +++ b/changelog.rst @@ -16,6 +16,7 @@ Bug fixes: * Add `VERSION` to built-in function completion so `SELECT VERSION();` is suggested. * Hide timezone notice at startup when local and server timezones are the same. * Let `sqlparse` accept arbitrarily-large queries. +* Respect user-specified `LIMIT` clauses when the limit value starts on a new line. 4.4.0 (2025-12-24) ================== diff --git a/pgcli/main.py b/pgcli/main.py index fe889dee6..ad01a46fd 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -27,6 +27,8 @@ from cli_helpers.utils import strip_ansi from .explain_output_formatter import ExplainOutputFormatter import click +import sqlparse +from sqlparse import tokens as sqlparse_tokens import tzlocal try: @@ -1114,7 +1116,7 @@ def _should_limit_output(self, sql, cur): def _has_limit(self, sql): if not sql: return False - return "limit " in sql.lower() + return any(token.match(sqlparse_tokens.Keyword, "LIMIT") for statement in sqlparse.parse(sql) for token in statement.flatten()) def _limit_output(self, cur): limit = min(self.row_limit, cur.rowcount) diff --git a/tests/test_rowlimit.py b/tests/test_rowlimit.py index da916b4da..9fe46ce06 100644 --- a/tests/test_rowlimit.py +++ b/tests/test_rowlimit.py @@ -44,9 +44,18 @@ def low_count(): return low_count_cursor -def test_row_limit_with_LIMIT_clause(LIMIT, over_limit): +@pytest.mark.parametrize( + "stmt", + [ + "SELECT * FROM students LIMIT 1000", + "SELECT * FROM students LIMIT\n 1000", + "SELECT * FROM students LIMIT\t1000", + "SELECT * FROM students limit 1000", + "SELECT * FROM students LiMiT 1000", + ], +) +def test_row_limit_with_LIMIT_clause(LIMIT, over_limit, stmt): cli = PGCli(row_limit=LIMIT) - stmt = "SELECT * FROM students LIMIT 1000" result = cli._should_limit_output(stmt, over_limit) assert result is False @@ -56,6 +65,21 @@ def test_row_limit_with_LIMIT_clause(LIMIT, over_limit): assert result is False +@pytest.mark.parametrize( + "stmt", + [ + "SELECT 'LIMIT 1000' FROM students", + "SELECT * FROM students -- LIMIT 1000", + "SELECT * FROM students /* LIMIT 1000 */", + ], +) +def test_row_limit_ignores_LIMIT_in_comments_or_strings(LIMIT, over_limit, stmt): + cli = PGCli(row_limit=LIMIT) + + result = cli._should_limit_output(stmt, over_limit) + assert result is True + + def test_row_limit_without_LIMIT_clause(LIMIT, over_limit): cli = PGCli(row_limit=LIMIT) stmt = "SELECT * FROM students"