[incubator-superset] 04/08: chore(sqla): assert query is single read-only statement (#11236)

[vi...@apache.org]

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

villebro pushed a commit to branch 0.38
in repository https://gitbox.apache.org/repos/asf/incubator-superset.git

commit 76f6e85eb3e82bac7fc9e27c78ae37c0cec2ebb2
Author: Ville Brofeldt <33...@users.noreply.github.com>
AuthorDate: Mon Oct 12 15:11:43 2020 +0300

    chore(sqla): assert query is single read-only statement (#11236)
---
 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 df07cca..934da17 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -64,6 +64,7 @@ from superset.jinja_context import (
 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 @@ class SqlaTable(  # pylint: disable=too-many-public-methods,too-many-instance-at
                     )
 
             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 9e18845..de46e1c 100644
--- a/tests/sqla_models_tests.py
+++ b/tests/sqla_models_tests.py
@@ -187,3 +187,45 @@ class TestDatabaseModel(SupersetTestCase):
         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)