feat: improve _extract_tables_from_sql (#26748)

This commit is contained in:
Beto Dealmeida
2024-03-18 13:02:58 -04:00
committed by Michael S. Molina
parent d4314a92d8
commit 94f677850c
4 changed files with 53 additions and 8 deletions

View File

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

View File

@@ -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 = []

View File

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