This is an automated email from the ASF dual-hosted git repository.
beto 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 bb002d6147 fix: is_select with UNION (#25290)
bb002d6147 is described below
commit bb002d6147c82de692f6692e77f59bd729953495
Author: Beto Dealmeida <[email protected]>
AuthorDate: Thu Sep 14 09:05:19 2023 -0700
fix: is_select with UNION (#25290)
---
superset/sql_parse.py | 10 +++++++---
tests/unit_tests/sql_parse_tests.py | 10 ++++++++++
2 files changed, 17 insertions(+), 3 deletions(-)
diff --git a/superset/sql_parse.py b/superset/sql_parse.py
index 34fc354730..d75551bef0 100644
--- a/superset/sql_parse.py
+++ b/superset/sql_parse.py
@@ -229,13 +229,17 @@ class ParsedQuery:
: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
diff --git a/tests/unit_tests/sql_parse_tests.py
b/tests/unit_tests/sql_parse_tests.py
index 73074d3df6..341ba9d789 100644
--- a/tests/unit_tests/sql_parse_tests.py
+++ b/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()
