From bb002d6147c82de692f6692e77f59bd729953495 Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Thu, 14 Sep 2023 09:05:19 -0700 Subject: [PATCH] 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 34fc354730536..d75551bef0767 100644 --- a/superset/sql_parse.py +++ b/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 diff --git a/tests/unit_tests/sql_parse_tests.py b/tests/unit_tests/sql_parse_tests.py index 73074d3df64bf..341ba9d789396 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()