fix: Ignore USE SQL keyword when determining SELECT statement (#28279)

This commit is contained in:
John Bodley 2024-05-02 11:25:55 -07:00 committed by GitHub
parent 0ce5864fc7
commit 27952e7057
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 2 deletions

View File

@ -884,6 +884,7 @@ class ParsedQuery:
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 @@ class ParsedQuery:
return False
if statement.get_type() == "SELECT":
seen_select = True
continue
if statement.get_type() != "UNKNOWN":
@ -920,8 +922,11 @@ class ParsedQuery:
):
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 @@ class ParsedQuery:
):
return False
return True
return seen_select
def get_inner_cte_expression(self, tokens: TokenList) -> TokenList | None:
for token in tokens:

View File

@ -1803,6 +1803,9 @@ WITH t AS (
)
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: