Skip to content

Commit

Permalink
fix: Ignore USE SQL keyword when determining SELECT statement (apache…
Browse files Browse the repository at this point in the history
  • Loading branch information
john-bodley authored and jzhao62 committed May 16, 2024
1 parent 9462344 commit 54c222f
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 2 deletions.
9 changes: 7 additions & 2 deletions superset/sql_parse.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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":
Expand All @@ -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(
Expand All @@ -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:
Expand Down
3 changes: 3 additions & 0 deletions tests/unit_tests/sql_parse_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit 54c222f

Please sign in to comment.