You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by be...@apache.org on 2023/05/11 23:30:45 UTC
[superset] 01/01: feat: robust(er) adhoc query validation
This is an automated email from the ASF dual-hosted git repository.
beto pushed a commit to branch robust_has_table_query
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 26feb0b61853e82625333f47e54dc18d43069d6a
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Thu May 11 15:37:29 2023 -0700
feat: robust(er) adhoc query validation
---
requirements/base.txt | 2 +
requirements/development.txt | 2 -
requirements/testing.txt | 2 +
setup.py | 1 +
superset/connectors/sqla/models.py | 27 +-------
superset/models/helpers.py | 69 ++++++++++++--------
superset/models/sql_lab.py | 3 +-
superset/sql_parse.py | 127 ++++++++++++++----------------------
tests/unit_tests/sql_parse_tests.py | 36 +++++-----
9 files changed, 119 insertions(+), 150 deletions(-)
diff --git a/requirements/base.txt b/requirements/base.txt
index 34478f1052..f4392b9c55 100644
--- a/requirements/base.txt
+++ b/requirements/base.txt
@@ -285,6 +285,8 @@ sqlalchemy-utils==0.38.3
# via
# apache-superset
# flask-appbuilder
+sqloxide==0.1.33
+ # via apache-superset
sqlparse==0.4.3
# via apache-superset
sshtunnel==0.4.0
diff --git a/requirements/development.txt b/requirements/development.txt
index 870c0c6a23..80f9449676 100644
--- a/requirements/development.txt
+++ b/requirements/development.txt
@@ -105,8 +105,6 @@ s3transfer==0.6.1
# via boto3
sasl==0.3.1
# via pyhive
-sqloxide==0.1.33
- # via -r requirements/development.in
stack-data==0.6.2
# via ipython
tableschema==1.20.2
diff --git a/requirements/testing.txt b/requirements/testing.txt
index 1e54713546..a4fd3f2731 100644
--- a/requirements/testing.txt
+++ b/requirements/testing.txt
@@ -107,6 +107,8 @@ pydata-google-auth==1.7.0
# via pandas-gbq
pyfakefs==5.2.2
# via -r requirements/testing.in
+pyhive[presto]==0.6.5
+ # via apache-superset
pytest==7.3.1
# via
# -r requirements/testing.in
diff --git a/setup.py b/setup.py
index b66b871818..22dd036d88 100644
--- a/setup.py
+++ b/setup.py
@@ -121,6 +121,7 @@ setup(
"slack_sdk>=3.19.0, <4",
"sqlalchemy>=1.4, <2",
"sqlalchemy-utils>=0.38.3, <0.39",
+ "sqloxide>=0.1",
"sqlparse>=0.4.3, <0.5",
"tabulate>=0.8.9, <0.9",
"typing-extensions>=4, <5",
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index 22d546424e..8a4a0f8aab 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -110,9 +110,9 @@ from superset.models.helpers import (
AuditMixinNullable,
CertificationMixin,
ExploreMixin,
+ process_sql_expression,
QueryResult,
QueryStringExtended,
- validate_adhoc_subquery,
)
from superset.sql_parse import ParsedQuery, sanitize_clause
from superset.superset_typing import AdhocColumn, AdhocMetric, Metric, QueryObjectDict
@@ -517,27 +517,6 @@ sqlatable_user = Table(
)
-def _process_sql_expression(
- expression: Optional[str],
- database_id: int,
- schema: str,
- template_processor: Optional[BaseTemplateProcessor] = None,
-) -> Optional[str]:
- if template_processor and expression:
- expression = template_processor.process_template(expression)
- if expression:
- try:
- expression = validate_adhoc_subquery(
- expression,
- database_id,
- schema,
- )
- expression = sanitize_clause(expression)
- except (QueryClauseValidationException, SupersetSecurityException) as ex:
- raise QueryObjectValidationError(ex.message) from ex
- return expression
-
-
class SqlaTable(
Model, BaseDatasource, ExploreMixin
): # pylint: disable=too-many-public-methods
@@ -999,7 +978,7 @@ class SqlaTable(
sqla_column = column(column_name)
sqla_metric = self.sqla_aggregations[metric["aggregate"]](sqla_column)
elif expression_type == utils.AdhocMetricExpressionType.SQL:
- expression = _process_sql_expression(
+ expression = process_sql_expression(
expression=metric["sqlExpression"],
database_id=self.database_id,
schema=self.schema,
@@ -1029,7 +1008,7 @@ class SqlaTable(
:rtype: sqlalchemy.sql.column
"""
label = utils.get_column_name(col)
- expression = _process_sql_expression(
+ expression = process_sql_expression(
expression=col["sqlExpression"],
database_id=self.database_id,
schema=self.schema,
diff --git a/superset/models/helpers.py b/superset/models/helpers.py
index 532b97bea6..8bd30a7003 100644
--- a/superset/models/helpers.py
+++ b/superset/models/helpers.py
@@ -114,9 +114,10 @@ ADVANCED_DATA_TYPES = config["ADVANCED_DATA_TYPES"]
def validate_adhoc_subquery(
- sql: str,
+ expression: str,
database_id: int,
default_schema: str,
+ sqla_dialect: Optional[str] = None,
) -> str:
"""
Check if adhoc SQL contains sub-queries or nested sub-queries with table.
@@ -128,9 +129,12 @@ def validate_adhoc_subquery(
:raise SupersetSecurityException if sql contains sub-queries or
nested sub-queries with table
"""
+ # build a proper SQL query from the expression
+ sql = f"SELECT {expression}"
+
statements = []
for statement in sqlparse.parse(sql):
- if has_table_query(statement):
+ if has_table_query(str(statement), sqla_dialect):
if not is_feature_enabled("ALLOW_ADHOC_SUBQUERY"):
raise SupersetSecurityException(
SupersetError(
@@ -145,6 +149,32 @@ def validate_adhoc_subquery(
return ";\n".join(str(statement) for statement in statements)
+def process_sql_expression(
+ expression: Optional[str],
+ database_id: int,
+ schema: str,
+ template_processor: Optional[BaseTemplateProcessor],
+ sqla_dialect: Optional[str] = None,
+) -> Optional[str]:
+ """
+ Apply template and validate ad-hoc expressions.
+ """
+ if template_processor and expression:
+ expression = template_processor.process_template(expression)
+ if expression:
+ expression = validate_adhoc_subquery(
+ expression,
+ database_id,
+ schema,
+ sqla_dialect,
+ )
+ try:
+ expression = sanitize_clause(expression)
+ except (QueryClauseValidationException, SupersetSecurityException) as ex:
+ raise QueryObjectValidationError(ex.message) from ex
+ return expression
+
+
def json_to_dict(json_str: str) -> Dict[Any, Any]:
if json_str:
val = re.sub(",[ \t\r\n]+}", "}", json_str)
@@ -838,27 +868,6 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
)
) from ex
- def _process_sql_expression( # pylint: disable=no-self-use
- self,
- expression: Optional[str],
- database_id: int,
- schema: str,
- template_processor: Optional[BaseTemplateProcessor],
- ) -> Optional[str]:
- if template_processor and expression:
- expression = template_processor.process_template(expression)
- if expression:
- expression = validate_adhoc_subquery(
- expression,
- database_id,
- schema,
- )
- try:
- expression = sanitize_clause(expression)
- except QueryClauseValidationException as ex:
- raise QueryObjectValidationError(ex.message) from ex
- return expression
-
def make_sqla_column_compatible(
self, sqla_col: ColumnElement, label: Optional[str] = None
) -> ColumnElement:
@@ -1140,11 +1149,12 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
sqla_column = sa.column(column_name)
sqla_metric = self.sqla_aggregations[metric["aggregate"]](sqla_column)
elif expression_type == utils.AdhocMetricExpressionType.SQL:
- expression = self._process_sql_expression(
+ expression = process_sql_expression(
expression=metric["sqlExpression"],
database_id=self.database_id,
schema=self.schema,
template_processor=template_processor,
+ sqla_dialect=self.db_engine_spec.engine,
)
sqla_metric = literal_column(expression)
else:
@@ -1523,11 +1533,12 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
if isinstance(col, dict):
col = cast(AdhocMetric, col)
if col.get("sqlExpression"):
- col["sqlExpression"] = self._process_sql_expression(
+ col["sqlExpression"] = process_sql_expression(
expression=col["sqlExpression"],
database_id=self.database_id,
schema=self.schema,
template_processor=template_processor,
+ sqla_dialect=self.db_engine_spec.engine,
)
if utils.is_adhoc_metric(col):
# add adhoc sort by column to columns_by_name if not exists
@@ -1589,6 +1600,7 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
selected,
self.database_id,
self.schema,
+ self.db_engine_spec.engine,
)
outer = literal_column(f"({selected})")
outer = self.make_sqla_column_compatible(outer, selected)
@@ -1615,6 +1627,7 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
_sql,
self.database_id,
self.schema,
+ self.db_engine_spec.engine,
)
select_exprs.append(
@@ -1876,11 +1889,12 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
msg=ex.message,
)
) from ex
- where = self._process_sql_expression(
+ where = process_sql_expression(
expression=where,
database_id=self.database_id,
schema=self.schema,
template_processor=template_processor,
+ sqla_dialect=self.db_engine_spec.engine,
)
where_clause_and += [self.text(where)]
having = extras.get("having")
@@ -1894,11 +1908,12 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
msg=ex.message,
)
) from ex
- having = self._process_sql_expression(
+ having = process_sql_expression(
expression=having,
database_id=self.database_id,
schema=self.schema,
template_processor=template_processor,
+ sqla_dialect=self.db_engine_spec.engine,
)
having_clause_and += [self.text(having)]
diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py
index b2f0c8c1ed..d840b5db0a 100644
--- a/superset/models/sql_lab.py
+++ b/superset/models/sql_lab.py
@@ -50,6 +50,7 @@ from superset.models.helpers import (
ExploreMixin,
ExtraJSONMixin,
ImportExportMixin,
+ process_sql_expression,
)
from superset.sql_parse import CtasMethod, ParsedQuery, Table
from superset.sqllab.limiting_factor import LimitingFactor
@@ -344,7 +345,7 @@ class Query(
:rtype: sqlalchemy.sql.column
"""
label = get_column_name(col)
- expression = self._process_sql_expression(
+ expression = process_sql_expression(
expression=col["sqlExpression"],
database_id=self.database_id,
schema=self.schema,
diff --git a/superset/sql_parse.py b/superset/sql_parse.py
index 323fabb3d3..7e889a4569 100644
--- a/superset/sql_parse.py
+++ b/superset/sql_parse.py
@@ -23,6 +23,7 @@ from urllib import parse
import sqlparse
from sqlalchemy import and_
+from sqloxide import parse_sql
from sqlparse.sql import (
Identifier,
IdentifierList,
@@ -47,16 +48,31 @@ from sqlparse.utils import imt
from superset.exceptions import QueryClauseValidationException
-try:
- from sqloxide import parse_sql as sqloxide_parse
-except: # pylint: disable=bare-except
- sqloxide_parse = None
-
RESULT_OPERATIONS = {"UNION", "INTERSECT", "EXCEPT", "SELECT"}
ON_KEYWORD = "ON"
PRECEDES_TABLE_NAME = {"FROM", "JOIN", "DESCRIBE", "WITH", "LEFT JOIN", "RIGHT JOIN"}
CTE_PREFIX = "CTE__"
+# mapping between sqloxide and SQLAlchemy dialects
+SQLOXIDE_DIALECTS = {
+ "ansi": {"trino", "trinonative", "presto"},
+ "hive": {"hive", "databricks"},
+ "ms": {"mssql"},
+ "mysql": {"mysql"},
+ "postgres": {
+ "cockroachdb",
+ "hana",
+ "netezza",
+ "postgres",
+ "postgresql",
+ "redshift",
+ "vertica",
+ },
+ "snowflake": {"snowflake"},
+ "sqlite": {"sqlite", "gsheets", "shillelagh"},
+ "clickhouse": {"clickhouse"},
+}
+
logger = logging.getLogger(__name__)
@@ -482,21 +498,11 @@ def sanitize_clause(clause: str) -> str:
return clause
-class InsertRLSState(str, Enum):
- """
- State machine that scans for WHERE and ON clauses referencing tables.
- """
-
- SCANNING = "SCANNING"
- SEEN_SOURCE = "SEEN_SOURCE"
- FOUND_TABLE = "FOUND_TABLE"
-
-
-def has_table_query(token_list: TokenList) -> bool:
+def has_table_query(sql: str, sqla_dialect: Optional[str] = None) -> bool:
"""
- Return if a statement has a query reading from a table.
+ Return if the sql has a subquery reading from a table.
- >>> has_table_query(sqlparse.parse("COUNT(*)")[0])
+ >>> has_table_query(sqlparse.parse("SELECT COUNT(*)")[0])
False
>>> has_table_query(sqlparse.parse("SELECT * FROM table")[0])
True
@@ -507,31 +513,17 @@ def has_table_query(token_list: TokenList) -> bool:
False
"""
- state = InsertRLSState.SCANNING
- for token in token_list.tokens:
- # Ignore comments
- if isinstance(token, sqlparse.sql.Comment):
- continue
-
- # Recurse into child token list
- if isinstance(token, TokenList) and has_table_query(token):
- return True
-
- # Found a source keyword (FROM/JOIN)
- if imt(token, m=[(Keyword, "FROM"), (Keyword, "JOIN")]):
- state = InsertRLSState.SEEN_SOURCE
+ return bool(extract_table_references(sql, sqla_dialect))
- # Found identifier/keyword after FROM/JOIN
- elif state == InsertRLSState.SEEN_SOURCE and (
- isinstance(token, sqlparse.sql.Identifier) or token.ttype == Keyword
- ):
- return True
- # Found nothing, leaving source
- elif state == InsertRLSState.SEEN_SOURCE and token.ttype != Whitespace:
- state = InsertRLSState.SCANNING
+class InsertRLSState(str, Enum):
+ """
+ State machine that scans for WHERE and ON clauses referencing tables.
+ """
- return False
+ SCANNING = "SCANNING"
+ SEEN_SOURCE = "SEEN_SOURCE"
+ FOUND_TABLE = "FOUND_TABLE"
def add_table_name(rls: TokenList, table: str) -> None:
@@ -711,57 +703,34 @@ def insert_rls(
return token_list
-# mapping between sqloxide and SQLAlchemy dialects
-SQLOXITE_DIALECTS = {
- "ansi": {"trino", "trinonative", "presto"},
- "hive": {"hive", "databricks"},
- "ms": {"mssql"},
- "mysql": {"mysql"},
- "postgres": {
- "cockroachdb",
- "hana",
- "netezza",
- "postgres",
- "postgresql",
- "redshift",
- "vertica",
- },
- "snowflake": {"snowflake"},
- "sqlite": {"sqlite", "gsheets", "shillelagh"},
- "clickhouse": {"clickhouse"},
-}
-
RE_JINJA_VAR = re.compile(r"\{\{[^\{\}]+\}\}")
RE_JINJA_BLOCK = re.compile(r"\{[%#][^\{\}%#]+[%#]\}")
def extract_table_references(
- sql_text: str, sqla_dialect: str, show_warning: bool = True
+ sql_text: str,
+ sqla_dialect: Optional[str] = None,
+ show_warning: bool = True,
) -> Set["Table"]:
"""
Return all the dependencies from a SQL sql_text.
"""
- dialect = "generic"
- tree = None
-
- if sqloxide_parse:
- for dialect, sqla_dialects in SQLOXITE_DIALECTS.items():
+ sqloxide_dialect = "generic"
+ if sqla_dialect:
+ for sqloxide_dialect, sqla_dialects in SQLOXIDE_DIALECTS.items():
if sqla_dialect in sqla_dialects:
break
- sql_text = RE_JINJA_BLOCK.sub(" ", sql_text)
- sql_text = RE_JINJA_VAR.sub("abc", sql_text)
- try:
- tree = sqloxide_parse(sql_text, dialect=dialect)
- except Exception as ex: # pylint: disable=broad-except
- if show_warning:
- logger.warning(
- "\nUnable to parse query with sqloxide:\n%s\n%s", sql_text, ex
- )
- # fallback to sqlparse
- if not tree:
- parsed = ParsedQuery(sql_text)
- return parsed.tables
+ sql_text = RE_JINJA_BLOCK.sub(" ", sql_text)
+ sql_text = RE_JINJA_VAR.sub("abc", sql_text)
+ try:
+ tree = parse_sql(sql_text, dialect=sqloxide_dialect)
+ except Exception as ex: # pylint: disable=broad-except
+ if show_warning:
+ logger.warning(
+ "\nUnable to parse query with sqloxide:\n%s\n%s", sql_text, ex
+ )
+ raise ex
def find_nodes_by_key(element: Any, target: str) -> Iterator[Any]:
"""
diff --git a/tests/unit_tests/sql_parse_tests.py b/tests/unit_tests/sql_parse_tests.py
index cfe6e213b2..351b301664 100644
--- a/tests/unit_tests/sql_parse_tests.py
+++ b/tests/unit_tests/sql_parse_tests.py
@@ -1206,16 +1206,16 @@ def test_sqlparse_issue_652():
@pytest.mark.parametrize(
"sql,expected",
[
- ("SELECT * FROM table", True),
- ("SELECT a FROM (SELECT 1 AS a) JOIN (SELECT * FROM table)", True),
+ ("SELECT * FROM 'table'", True),
+ ("SELECT a FROM (SELECT 1 AS a) JOIN (SELECT * FROM 'table')", True),
("(SELECT COUNT(DISTINCT name) AS foo FROM birth_names)", True),
- ("COUNT(*)", False),
+ ("SELECT COUNT(*)", False),
("SELECT a FROM (SELECT 1 AS a)", False),
- ("SELECT a FROM (SELECT 1 AS a) JOIN table", True),
+ ("SELECT a FROM (SELECT 1 AS a) JOIN 'table'", True),
("SELECT * FROM (SELECT 1 AS foo, 2 AS bar) ORDER BY foo ASC, bar", False),
("SELECT * FROM other_table", True),
- ("extract(HOUR from from_unixtime(hour_ts)", False),
- ("(SELECT * FROM table)", True),
+ ("SELECT extract(HOUR from from_unixtime(hour_ts))", False),
+ ("(SELECT * FROM 'table')", True),
("(SELECT COUNT(DISTINCT name) from birth_names)", True),
(
"(SELECT table_name FROM information_schema.tables WHERE table_name LIKE '%user%' LIMIT 1)",
@@ -1225,6 +1225,10 @@ def test_sqlparse_issue_652():
"(SELECT table_name FROM /**/ information_schema.tables WHERE table_name LIKE '%user%' LIMIT 1)",
True,
),
+ (
+ "SELECT table_name FROM (information_schema.tables) WHERE table_name LIKE '%user%' LIMIT 1",
+ True,
+ ),
],
)
def test_has_table_query(sql: str, expected: bool) -> None:
@@ -1234,8 +1238,7 @@ def test_has_table_query(sql: str, expected: bool) -> None:
This is used to prevent ad-hoc metrics from querying unauthorized tables, bypassing
row-level security.
"""
- statement = sqlparse.parse(sql)[0]
- assert has_table_query(statement) == expected
+ assert has_table_query(sql) == expected
@pytest.mark.parametrize(
@@ -1521,18 +1524,17 @@ def test_extract_table_references(mocker: MockerFixture) -> None:
Table(table="other_table", schema=None, catalog=None),
}
- # test falling back to sqlparse
- logger = mocker.patch("superset.sql_parse.logger")
- sql = "SELECT * FROM table UNION ALL SELECT * FROM other_table"
+ sql = "SELECT * FROM 'table' UNION ALL SELECT * FROM other_table"
assert extract_table_references(
sql,
"trino",
- ) == {Table(table="other_table", schema=None, catalog=None)}
- logger.warning.assert_called_once()
+ ) == {
+ Table(table="table", schema=None, catalog=None),
+ Table(table="other_table", schema=None, catalog=None),
+ }
- logger = mocker.patch("superset.migrations.shared.utils.logger")
- sql = "SELECT * FROM table UNION ALL SELECT * FROM other_table"
+ sql = "SELECT * FROM 'table' UNION ALL SELECT * FROM other_table"
assert extract_table_references(sql, "trino", show_warning=False) == {
- Table(table="other_table", schema=None, catalog=None)
+ Table(table="table", schema=None, catalog=None),
+ Table(table="other_table", schema=None, catalog=None),
}
- logger.warning.assert_not_called()