diff --git a/superset/sql_parse.py b/superset/sql_parse.py index 034e2083ffc7c..f32647042b0aa 100644 --- a/superset/sql_parse.py +++ b/superset/sql_parse.py @@ -884,6 +884,7 @@ def is_body_select(body: dict[str, Any]) -> bool: def is_select(self) -> bool: # make sure we strip comments; prevents a bug with comments in the CTE parsed = sqlparse.parse(self.strip_comments()) + seen_select = False for statement in parsed: # Check if this is a CTE @@ -907,6 +908,7 @@ def is_select(self) -> bool: return False if statement.get_type() == "SELECT": + seen_select = True continue if statement.get_type() != "UNKNOWN": @@ -920,8 +922,11 @@ def is_select(self) -> bool: ): return False + if imt(statement.tokens[0], m=(Keyword, "USE")): + continue + # return false on `EXPLAIN`, `SET`, `SHOW`, etc. - if statement[0].ttype == Keyword: + if imt(statement.tokens[0], t=Keyword): return False if not any( @@ -930,7 +935,7 @@ def is_select(self) -> bool: ): return False - return True + return seen_select def get_inner_cte_expression(self, tokens: TokenList) -> TokenList | None: for token in tokens: diff --git a/tests/unit_tests/sql_parse_tests.py b/tests/unit_tests/sql_parse_tests.py index 3a78885abec07..3b80b7e01d1ac 100644 --- a/tests/unit_tests/sql_parse_tests.py +++ b/tests/unit_tests/sql_parse_tests.py @@ -1803,6 +1803,9 @@ def test_is_select() -> None: ) SELECT * FROM t""" ).is_select() + assert not ParsedQuery("").is_select() + assert not ParsedQuery("USE foo").is_select() + assert ParsedQuery("USE foo; SELECT * FROM bar").is_select() def test_sqlquery() -> None: