fix: is_select with UNION (apache#25290)

Author: betodealmeida

superset/sql_parse.py

@@ -229,13 +229,17 @@ def _check_cte_is_select(self, oxide_parse: list[dict[str, Any]]) -> bool:
         :param oxide_parse: parsed CTE
         :return: True if CTE is a SELECT statement
         """
+
+        def is_body_select(body: dict[str, Any]) -> bool:
+            if op := body.get("SetOperation"):
+                return is_body_select(op["left"]) and is_body_select(op["right"])
+            return all(key == "Select" for key in body.keys())
+
         for query in oxide_parse:
             parsed_query = query["Query"]
             cte_tables = self._get_cte_tables(parsed_query)
             for cte_table in cte_tables:
-                is_select = all(
-                    key == "Select" for key in cte_table["query"]["body"].keys()
-                )
+                is_select = is_body_select(cte_table["query"]["body"])
                 if not is_select:
                     return False
         return True

tests/unit_tests/sql_parse_tests.py

@@ -1623,3 +1623,13 @@ def test_is_select() -> None:
     Test `is_select`.
     """
     assert not ParsedQuery("SELECT 1; DROP DATABASE superset").is_select()
+    assert ParsedQuery(
+        "with base as(select id from table1 union all select id from table2) select * from base"
+    ).is_select()
+    assert ParsedQuery(
+        """
+WITH t AS (
+    SELECT 1 UNION ALL SELECT 2
+)
+SELECT * FROM t"""
+    ).is_select()