Skip to content

Commit

Permalink
fix: is_select with UNION (apache#25290)
Browse files Browse the repository at this point in the history
  • Loading branch information
betodealmeida authored Sep 14, 2023
1 parent e1ddba9 commit bb002d6
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 3 deletions.
10 changes: 7 additions & 3 deletions superset/sql_parse.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 10 additions & 0 deletions tests/unit_tests/sql_parse_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

0 comments on commit bb002d6

Please sign in to comment.