feat: improve logic in is_select (#17329)

* feat: improve logic in is_select

* Add more edge cases

tests/unit_tests/sql_parse_tests.py

@@ -15,10 +15,17 @@
 # specific language governing permissions and limitations
 # under the License.
 
+# pylint: disable=invalid-name
+
+import sqlparse
+
 from superset.sql_parse import ParsedQuery
 
 
-def test_cte_with_comments():
+def test_cte_with_comments_is_select():
+    """
+    Some CTES with comments are not correctly identified as SELECTS.
+    """
     sql = ParsedQuery(
         """WITH blah AS
   (SELECT * FROM core_dev.manager_team),
@@ -44,3 +51,43 @@ SELECT * FROM blah
 INNER JOIN blah2 ON blah2.team_id = blah.team_id"""
     )
     assert sql.is_select()
+
+
+def test_cte_is_select():
+    """
+    Some CTEs are not correctly identified as SELECTS.
+    """
+    # `AS(` gets parsed as a function
+    sql = ParsedQuery(
+        """WITH foo AS(
+SELECT
+  FLOOR(__time TO WEEK) AS "week",
+  name,
+  COUNT(DISTINCT user_id) AS "unique_users"
+FROM "druid"."my_table"
+GROUP BY 1,2
+)
+SELECT
+  f.week,
+  f.name,
+  f.unique_users
+FROM foo f"""
+    )
+    assert sql.is_select()
+
+
+def test_unknown_select():
+    """
+    Test that `is_select` works when sqlparse fails to identify the type.
+    """
+    sql = "WITH foo AS(SELECT 1) SELECT 1"
+    assert sqlparse.parse(sql)[0].get_type() == "UNKNOWN"
+    assert ParsedQuery(sql).is_select()
+
+    sql = "WITH foo AS(SELECT 1) INSERT INTO my_table (a) VALUES (1)"
+    assert sqlparse.parse(sql)[0].get_type() == "UNKNOWN"
+    assert not ParsedQuery(sql).is_select()
+
+    sql = "WITH foo AS(SELECT 1) DELETE FROM my_table"
+    assert sqlparse.parse(sql)[0].get_type() == "UNKNOWN"
+    assert not ParsedQuery(sql).is_select()