From 94827665de50029f71905d71aefaa0a71087fb7b Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Mon, 12 Oct 2020 14:09:10 +0300 Subject: [PATCH] chore(sqla): assert query is single read-only statement --- superset/connectors/sqla/models.py | 9 +++++++ tests/sqla_models_tests.py | 42 ++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index df07ccad67a89..934da17c73dc7 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -64,6 +64,7 @@ from superset.models.annotations import Annotation from superset.models.core import Database from superset.models.helpers import AuditMixinNullable, QueryResult +from superset.sql_parse import ParsedQuery from superset.typing import Metric, QueryObjectDict from superset.utils import core as utils, import_datasource @@ -755,6 +756,14 @@ def get_from_clause( ) from_sql = sqlparse.format(from_sql, strip_comments=True) + if len(sqlparse.split(from_sql)) > 1: + raise QueryObjectValidationError( + _("Virtual dataset query cannot consist of multiple statements") + ) + if not ParsedQuery(from_sql).is_readonly(): + raise QueryObjectValidationError( + _("Virtual dataset query must be read-only") + ) return TextAsFrom(sa.text(from_sql), []).alias("expr_qry") return self.get_sqla_table() diff --git a/tests/sqla_models_tests.py b/tests/sqla_models_tests.py index 9e1884503d4c4..de46e1c555a4f 100644 --- a/tests/sqla_models_tests.py +++ b/tests/sqla_models_tests.py @@ -187,3 +187,45 @@ def test_incorrect_jinja_syntax_raises_correct_exception(self): if get_example_database().backend != "presto": with pytest.raises(QueryObjectValidationError): table.get_sqla_query(**query_obj) + + def test_multiple_sql_statements_raises_exception(self): + base_query_obj = { + "granularity": None, + "from_dttm": None, + "to_dttm": None, + "groupby": ["grp"], + "metrics": [], + "is_timeseries": False, + "filter": [], + } + + table = SqlaTable( + table_name="test_has_extra_cache_keys_table", + sql="SELECT 'foo' as grp, 1 as num; SELECT 'bar' as grp, 2 as num", + database=get_example_database(), + ) + + query_obj = dict(**base_query_obj, extras={}) + with pytest.raises(QueryObjectValidationError): + table.get_sqla_query(**query_obj) + + def test_dml_statement_raises_exception(self): + base_query_obj = { + "granularity": None, + "from_dttm": None, + "to_dttm": None, + "groupby": ["grp"], + "metrics": [], + "is_timeseries": False, + "filter": [], + } + + table = SqlaTable( + table_name="test_has_extra_cache_keys_table", + sql="DELETE FROM foo", + database=get_example_database(), + ) + + query_obj = dict(**base_query_obj, extras={}) + with pytest.raises(QueryObjectValidationError): + table.get_sqla_query(**query_obj)