(superset) 08/09: feat: `improve _extract_tables_from_sql` (#26748)

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

michaelsmolina pushed a commit to branch 3.1
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 8274d869b1d5a7bbe1992b9dc91b7c8d233f4740
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Mon Mar 18 13:02:58 2024 -0400

    feat: `improve _extract_tables_from_sql` (#26748)
---
 superset/sql_parse.py                     | 18 +++++++++++++---
 tests/unit_tests/jinja_context_test.py    |  8 ++++++++
 tests/unit_tests/security/manager_test.py |  1 +
 tests/unit_tests/sql_parse_tests.py       | 34 ++++++++++++++++++++++++++-----
 4 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/superset/sql_parse.py b/superset/sql_parse.py
index c85afc9460..49ee4c491f 100644
--- a/superset/sql_parse.py
+++ b/superset/sql_parse.py
@@ -25,6 +25,7 @@ from dataclasses import dataclass
 from typing import Any, cast, Optional
 
 import sqlparse
+from flask_babel import gettext as __
 from sqlalchemy import and_
 from sqlglot import exp, parse, parse_one
 from sqlglot.dialects import Dialects
@@ -55,7 +56,11 @@ from sqlparse.tokens import (
 )
 from sqlparse.utils import imt
 
-from superset.exceptions import QueryClauseValidationException
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import (
+    QueryClauseValidationException,
+    SupersetSecurityException,
+)
 from superset.utils.backports import StrEnum
 
 try:
@@ -287,9 +292,16 @@ class ParsedQuery:
         """
         try:
             statements = parse(self.stripped(), dialect=self._dialect)
-        except SqlglotError:
+        except SqlglotError as ex:
             logger.warning("Unable to parse SQL (%s): %s", self._dialect, self.sql)
-            return set()
+            dialect = self._dialect or "generic"
+            raise SupersetSecurityException(
+                SupersetError(
+                    error_type=SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR,
+                    message=__(f"Unable to parse SQL ({dialect}): {self.sql}"),
+                    level=ErrorLevel.ERROR,
+                )
+            ) from ex
 
         return {
             table
diff --git a/tests/unit_tests/jinja_context_test.py b/tests/unit_tests/jinja_context_test.py
index e2a5e8cd49..bbd9bcf06c 100644
--- a/tests/unit_tests/jinja_context_test.py
+++ b/tests/unit_tests/jinja_context_test.py
@@ -90,6 +90,14 @@ def test_dataset_macro(mocker: MockFixture) -> None:
     )
     DatasetDAO = mocker.patch("superset.daos.dataset.DatasetDAO")
     DatasetDAO.find_by_id.return_value = dataset
+    mocker.patch(
+        "superset.connectors.sqla.models.security_manager.get_guest_rls_filters",
+        return_value=[],
+    )
+    mocker.patch(
+        "superset.models.helpers.security_manager.get_guest_rls_filters",
+        return_value=[],
+    )
 
     assert (
         dataset_macro(1)
diff --git a/tests/unit_tests/security/manager_test.py b/tests/unit_tests/security/manager_test.py
index 5a06013a68..325531c25b 100644
--- a/tests/unit_tests/security/manager_test.py
+++ b/tests/unit_tests/security/manager_test.py
@@ -202,6 +202,7 @@ def test_raise_for_access_query_default_schema(
     sm = SupersetSecurityManager(appbuilder)
     mocker.patch.object(sm, "can_access_database", return_value=False)
     mocker.patch.object(sm, "get_schema_perm", return_value="[PostgreSQL].[public]")
+    mocker.patch.object(sm, "is_guest_user", return_value=False)
     SqlaTable = mocker.patch("superset.connectors.sqla.models.SqlaTable")
     SqlaTable.query_datasources_by_name.return_value = []
 
diff --git a/tests/unit_tests/sql_parse_tests.py b/tests/unit_tests/sql_parse_tests.py
index 2d2448c2fa..fc1ae1b231 100644
--- a/tests/unit_tests/sql_parse_tests.py
+++ b/tests/unit_tests/sql_parse_tests.py
@@ -25,7 +25,10 @@ from sqlalchemy import text
 from sqlparse.sql import Identifier, Token, TokenList
 from sqlparse.tokens import Name
 
-from superset.exceptions import QueryClauseValidationException
+from superset.exceptions import (
+    QueryClauseValidationException,
+    SupersetSecurityException,
+)
 from superset.sql_parse import (
     add_table_name,
     extract_table_references,
@@ -265,13 +268,34 @@ def test_extract_tables_illdefined() -> None:
     """
     Test that ill-defined tables return an empty set.
     """
-    assert extract_tables("SELECT * FROM schemaname.") == set()
-    assert extract_tables("SELECT * FROM catalogname.schemaname.") == set()
-    assert extract_tables("SELECT * FROM catalogname..") == set()
+    with pytest.raises(SupersetSecurityException) as excinfo:
+        extract_tables("SELECT * FROM schemaname.")
+    assert (
+        str(excinfo.value) == "Unable to parse SQL (generic): SELECT * FROM schemaname."
+    )
+
+    with pytest.raises(SupersetSecurityException) as excinfo:
+        extract_tables("SELECT * FROM catalogname.schemaname.")
+    assert (
+        str(excinfo.value)
+        == "Unable to parse SQL (generic): SELECT * FROM catalogname.schemaname."
+    )
+
+    with pytest.raises(SupersetSecurityException) as excinfo:
+        extract_tables("SELECT * FROM catalogname..")
+    assert (
+        str(excinfo.value)
+        == "Unable to parse SQL (generic): SELECT * FROM catalogname.."
+    )
+
+    with pytest.raises(SupersetSecurityException) as excinfo:
+        extract_tables('SELECT * FROM "tbname')
+    assert str(excinfo.value) == 'Unable to parse SQL (generic): SELECT * FROM "tbname'
+
+    # odd edge case that works
     assert extract_tables("SELECT * FROM catalogname..tbname") == {
         Table(table="tbname", schema=None, catalog="catalogname")
     }
-    assert extract_tables('SELECT * FROM "tbname') == set()
 
 
 def test_extract_tables_show_tables_from() -> None: