iop
/
superset
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: is_select with UNION (#25290)

This commit is contained in:
Beto Dealmeida 2023-09-14 09:05:19 -07:00 committed by GitHub
parent e1ddba9c0f
commit bb002d6147
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 17 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
superset/sql_parse.py
View File

@ -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

10
tests/unit_tests/sql_parse_tests.py
View File

@ -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()