diff --git a/AUTHORS b/AUTHORS index 06c8d9a17..853110bfb 100644 --- a/AUTHORS +++ b/AUTHORS @@ -152,6 +152,7 @@ Contributors: * Shayan Golshani (shgol) * Tommi Kyntölä (kynde) * Diego + * VXNCXNX Creator: -------- diff --git a/changelog.rst b/changelog.rst index 0c3905e38..f45af4ba8 100644 --- a/changelog.rst +++ b/changelog.rst @@ -7,6 +7,8 @@ Bug fixes: * Fix ``TypeError: cannot use a string pattern on a bytes-like object`` when completion metadata comes back as bytes (e.g. ``SQL_ASCII`` client encoding). * Suggest columns, not datatypes, after a column literally named ``type`` in a ``SELECT`` list. +* Detect an unconditional ``UPDATE`` with ``sqlparse`` rather than splitting on whitespace, so a + ``WHERE`` appearing inside a string literal no longer suppresses the destructive-statement warning. Features: --------- diff --git a/pgcli/packages/parseutils/__init__.py b/pgcli/packages/parseutils/__init__.py index 434a4cc73..daecc189a 100644 --- a/pgcli/packages/parseutils/__init__.py +++ b/pgcli/packages/parseutils/__init__.py @@ -20,10 +20,25 @@ def query_starts_with(formatted_sql, prefixes): return bool(formatted_sql) and formatted_sql.split()[0] in prefixes -def query_is_unconditional_update(formatted_sql): - """Check if the query starts with UPDATE and contains no WHERE.""" - tokens = formatted_sql.split() - return bool(tokens) and tokens[0] == "update" and "where" not in tokens +def query_is_unconditional_update(query): + """Check if the query starts with UPDATE and contains no top-level WHERE clause. + + Uses sqlparse's parse tree (rather than naive whitespace splitting) so that + the word "where" appearing inside a string literal, comment, or a nested + subquery doesn't get mistaken for an actual WHERE clause. + """ + statements = sqlparse.parse(query) + if not statements: + return False + statement = statements[0] + + first_token = statement.token_first(skip_cm=True) + if first_token is None or first_token.ttype is not sqlparse.tokens.DML: + return False + if first_token.value.upper() != "UPDATE": + return False + + return not any(isinstance(token, sqlparse.sql.Where) for token in statement.tokens) def is_destructive(queries, keywords): @@ -31,7 +46,7 @@ def is_destructive(queries, keywords): for query in sqlparse.split(queries): if query: formatted_sql = sqlparse.format(query.lower(), strip_comments=True).strip() - if "unconditional_update" in keywords and query_is_unconditional_update(formatted_sql): + if "unconditional_update" in keywords and query_is_unconditional_update(query): return True if query_starts_with(formatted_sql, keywords): return True diff --git a/tests/parseutils/test_parseutils.py b/tests/parseutils/test_parseutils.py index 90749ebfc..02341125f 100644 --- a/tests/parseutils/test_parseutils.py +++ b/tests/parseutils/test_parseutils.py @@ -285,6 +285,28 @@ def test_is_destructive(sql, keywords, expected): assert is_destructive(sql, keywords) == expected +@pytest.mark.parametrize( + ("sql", "expected"), + [ + # A "where" appearing inside a string literal must not be mistaken + # for a WHERE clause (regression test for the unconditional UPDATE + # confirmation bypass). + ("update accounts set note = 'no where clause here'", True), + ("update t set c = 'nowhere'", True), + ("update accounts set balance = 0", True), + ("update accounts set balance = 0 where id = 1", False), + ("UPDATE t SET c = 1", True), + ("-- where\nupdate t set c = 1", True), + ("select * from t", False), + ("", False), + # A WHERE inside a subquery does not make the outer UPDATE conditional. + ("update t set c = (select x from y where z = 1)", True), + ], +) +def test_is_destructive_unconditional_update_string_literal(sql, expected): + assert is_destructive(sql, ["unconditional_update"]) == expected + + @pytest.mark.parametrize( ("warning_level", "expected"), [