This is an automated email from the ASF dual-hosted git repository. michaelsmolina pushed a commit to branch 3.0 in repository https://gitbox.apache.org/repos/asf/superset.git
commit 807a027a5ffb27c4335cc8af81dd31d24559e165 Author: Beto Dealmeida <[email protected]> AuthorDate: Thu Sep 14 09:05:19 2023 -0700 fix: is_select with UNION (#25290) (cherry picked from commit bb002d6147c82de692f6692e77f59bd729953495) --- 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 2d467ec2b3..c196fdabfa 100644 --- a/superset/sql_parse.py +++ b/superset/sql_parse.py @@ -230,13 +230,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()
