diff --git a/superset/sql/parse.py b/superset/sql/parse.py index c9ed22ff0f5..bcc6b32382b 100644 --- a/superset/sql/parse.py +++ b/superset/sql/parse.py @@ -654,6 +654,16 @@ class SQLStatement(BaseSQLStatement[exp.Expression]): if isinstance(self._parsed, exp.Command) and self._parsed.name == "ALTER": return True # pragma: no cover + if ( + self._dialect == Dialects.POSTGRES + and isinstance(self._parsed, exp.Command) + and self._parsed.name == "DO" + ): + # anonymous blocks can be written in many different languages (the default + # is PL/pgSQL), so parsing them it out of scope of this class; we just + # assume the anonymous block is mutating + return True + # Postgres runs DMLs prefixed by `EXPLAIN ANALYZE`, see # https://www.postgresql.org/docs/current/sql-explain.html if ( diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py index 458caf5fa11..6e4db3c8880 100644 --- a/tests/unit_tests/sql/parse_tests.py +++ b/tests/unit_tests/sql/parse_tests.py @@ -1189,6 +1189,43 @@ def test_is_mutating(sql: str, engine: str, expected: bool) -> None: assert SQLStatement(sql, engine).is_mutating() == expected +@pytest.mark.parametrize( + "sql, expected", + [ + ( + """ +DO $$ +BEGIN + INSERT INTO public.users (name, real_name) + VALUES ('SQLLab bypass DML', 'SQLLab bypass DML'); +END; +$$; + """, + True, + ), + ( + """ +DO $$ +BEGIN + IF (SELECT COUNT(*) FROM orders WHERE status = 'pending') > 100 THEN + RAISE NOTICE 'High pending order volume detected'; + END IF; +END; +$$; + """, + True, + ), + ], +) +def test_is_mutating_anonymous_block(sql: str, expected: bool) -> None: + """ + Test for `is_mutating` with a Postgres anonymous block. + + Since we can't parse the PL/pgSQL inside the block we always assume it is mutating. + """ + assert SQLStatement(sql, "postgresql").is_mutating() == expected + + def test_optimize() -> None: """ Test that the `optimize` method works as expected.