From 54c222fede3e7bff90a5f943de39996e49febc2f Mon Sep 17 00:00:00 2001
From: John Bodley <4567245+john-bodley@users.noreply.github.com>
Date: Thu, 2 May 2024 11:25:55 -0700
Subject: [PATCH] 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 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: