(superset) branch master updated: fix: Ignore USE SQL keyword when determining SELECT statement (#28279)

[johnbodley]

This is an automated email from the ASF dual-hosted git repository.

johnbodley pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new 27952e7057 fix: Ignore USE SQL keyword when determining SELECT 
statement (#28279)
27952e7057 is described below

commit 27952e705754802fa5127e467012607bd892cef2
Author: John Bodley <4567245+john-bod...@users.noreply.github.com>
AuthorDate: Thu May 2 11:25:55 2024 -0700

    fix: Ignore USE SQL keyword when determining SELECT statement (#28279)
---
 superset/sql_parse.py               | 9 +++++++--
 tests/unit_tests/sql_parse_tests.py | 3 +++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/superset/sql_parse.py b/superset/sql_parse.py
index 034e2083ff..f32647042b 100644
--- a/superset/sql_parse.py
+++ b/superset/sql_parse.py
@@ -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:
diff --git a/tests/unit_tests/sql_parse_tests.py 
b/tests/unit_tests/sql_parse_tests.py
index 3a78885abe..3b80b7e01d 100644
--- a/tests/unit_tests/sql_parse_tests.py
+++ b/tests/unit_tests/sql_parse_tests.py
@@ -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: