diff --git a/superset/sql/parse.py b/superset/sql/parse.py index c750e0551b01..14830c9e876b 100644 --- a/superset/sql/parse.py +++ b/superset/sql/parse.py @@ -1373,6 +1373,22 @@ def extract_tables_from_statement( if isinstance(source, exp.Table) and not is_cte(source, scope) ] + # Use of `traverse_scope` doesn't always catch tables in subqueries within + # `VALUES` clauses, so we need to search for them explicitly. + # See https://github.com/apache/superset/issues/31599 for more details. + cte_names = { + cte.alias + for cte in statement.find_all(exp.CTE) + if isinstance(cte.alias, str) + } + for value in statement.find_all(exp.Values): + for table in value.find_all(exp.Table): + # If the table is not a CTE, it should be treated as a source. + # This is a heuristic, as we can't easily check for CTEs in scopes + # without traversing them. + if table.name not in cte_names: + sources.append(table) + return { Table( source.name, diff --git a/tests/unit_tests/sql/parse_tests_vulnerability.py b/tests/unit_tests/sql/parse_tests_vulnerability.py new file mode 100644 index 000000000000..91437c8edfc3 --- /dev/null +++ b/tests/unit_tests/sql/parse_tests_vulnerability.py @@ -0,0 +1,41 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import unittest +from superset.sql.parse import SQLStatement + +class TestSqlParseVulnerability(unittest.TestCase): + def test_values_subquery_extraction(self): + """ + Test that tables hidden inside a VALUES clause with a nested subquery are extracted. + """ + sql = "SUM(revenue) + (SELECT \"r\" FROM (VALUES((SELECT COUNT(*) FROM sales WHERE country='CA'))) v(\"r\"))" + statement = SQLStatement(sql, engine="postgresql") + tables = statement.tables + + self.assertIn("sales", [t.table for t in tables], "Should extract 'sales' table from nested VALUES subquery") + + def test_normal_subquery_extraction(self): + """ + Test that normal subqueries work as expected. + """ + sql = "SELECT * FROM (SELECT * FROM sales) AS sub" + statement = SQLStatement(sql, engine="postgresql") + tables = statement.tables + self.assertIn("sales", [t.table for t in tables]) + +if __name__ == '__main__': + unittest.main()