chore: Support SET & SHOW commands as read only SQL commands (apache#…

…11868)

* Support SET & SHOW commands as read only SQL commands

* Move is_readonly definition into the engine spec

* Rename & use super()

Co-authored-by: bogdan kyryliuk <[email protected]>

superset/connectors/sqla/models.py

@@ -627,7 +627,7 @@ def external_metadata(self) -> List[Dict[str, str]]:
             engine = self.database.get_sqla_engine(schema=self.schema)
             sql = self.get_template_processor().process_template(self.sql)
             parsed_query = ParsedQuery(sql)
-            if not parsed_query.is_readonly():
+            if not db_engine_spec.is_readonly_query(parsed_query):
                 raise SupersetSecurityException(
                     SupersetError(
                         error_type=SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR,
@@ -799,7 +799,11 @@ def get_from_clause(
                     _("Virtual dataset query cannot consist of multiple statements")
                 )
             parsed_query = ParsedQuery(from_sql)
-            if not (parsed_query.is_unknown() or parsed_query.is_readonly()):
+            db_engine_spec = self.database.db_engine_spec
+            if not (
+                parsed_query.is_unknown()
+                or db_engine_spec.is_readonly_query(parsed_query)
+            ):
                 raise QueryObjectValidationError(
                     _("Virtual dataset query must be read-only")
                 )

superset/db_engine_specs/base.py

@@ -54,7 +54,7 @@
 from superset import app, sql_parse
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.models.sql_lab import Query
-from superset.sql_parse import Table
+from superset.sql_parse import ParsedQuery, Table
 from superset.utils import core as utils
 
 if TYPE_CHECKING:
@@ -1038,3 +1038,8 @@ def get_extra_params(database: "Database") -> Dict[str, Any]:
                 logger.error(ex)
                 raise ex
         return extra
+
+    @classmethod
+    def is_readonly_query(cls, parsed_query: ParsedQuery) -> bool:
+        """Pessimistic readonly, 100% sure statement won't mutate anything"""
+        return parsed_query.is_select() or parsed_query.is_explain()

superset/db_engine_specs/hive.py

@@ -37,13 +37,14 @@
 from superset.exceptions import SupersetException
 from superset.extensions import cache_manager
 from superset.models.sql_lab import Query
-from superset.sql_parse import Table
+from superset.sql_parse import ParsedQuery, Table
 from superset.utils import core as utils
 
 if TYPE_CHECKING:
     # prevent circular imports
     from superset.models.core import Database
 
+
 QueryStatus = utils.QueryStatus
 config = app.config
 logger = logging.getLogger(__name__)
@@ -525,3 +526,12 @@ def get_function_names(cls, database: "Database") -> List[str]:
         :return: A list of function names useable in the database
         """
         return database.get_df("SHOW FUNCTIONS")["tab_name"].tolist()
+
+    @classmethod
+    def is_readonly_query(cls, parsed_query: ParsedQuery) -> bool:
+        """Pessimistic readonly, 100% sure statement won't mutate anything"""
+        return (
+            super().is_readonly_query(parsed_query)
+            or parsed_query.is_set()
+            or parsed_query.is_show()
+        )

superset/db_engine_specs/presto.py

@@ -111,7 +111,7 @@ def get_children(column: Dict[str, str]) -> List[Dict[str, str]]:
     raise Exception(f"Unknown type {type_}!")
 
 
-class PrestoEngineSpec(BaseEngineSpec):
+class PrestoEngineSpec(BaseEngineSpec):  # pylint: disable=too-many-public-methods
     engine = "presto"
     engine_name = "Presto"
 
@@ -1090,3 +1090,8 @@ def extract_errors(cls, ex: Exception) -> List[Dict[str, Any]]:
                 )
             )
         ]
+
+    @classmethod
+    def is_readonly_query(cls, parsed_query: ParsedQuery) -> bool:
+        """Pessimistic readonly, 100% sure statement won't mutate anything"""
+        return super().is_readonly_query(parsed_query) or parsed_query.is_show()

superset/sql_lab.py

@@ -167,7 +167,7 @@ def execute_sql_statement(
     parsed_query = ParsedQuery(sql_statement)
     sql = parsed_query.stripped()
 
-    if not parsed_query.is_readonly() and not database.allow_dml:
+    if not db_engine_spec.is_readonly_query(parsed_query) and not database.allow_dml:
         raise SqlLabSecurityException(
             _("Only `SELECT` statements are allowed against this database")
         )

superset/sql_parse.py

@@ -119,13 +119,25 @@ def is_explain(self) -> bool:
         # Explain statements will only be the first statement
         return statements_without_comments.startswith("EXPLAIN")
 
+    def is_show(self) -> bool:
+        # Remove comments
+        statements_without_comments = sqlparse.format(
+            self.stripped(), strip_comments=True
+        )
+        # Show statements will only be the first statement
+        return statements_without_comments.upper().startswith("SHOW")
+
+    def is_set(self) -> bool:
+        # Remove comments
+        statements_without_comments = sqlparse.format(
+            self.stripped(), strip_comments=True
+        )
+        # Set statements will only be the first statement
+        return statements_without_comments.upper().startswith("SET")
+
     def is_unknown(self) -> bool:
         return self._parsed[0].get_type() == "UNKNOWN"
 
-    def is_readonly(self) -> bool:
-        """Pessimistic readonly, 100% sure statement won't mutate anything"""
-        return self.is_select() or self.is_explain()
-
     def stripped(self) -> str:
         return self.sql.strip(" \t\n;")
 

tests/core_tests.py

@@ -36,6 +36,7 @@
 import sqlalchemy as sqla
 
 from superset.models.cache import CacheKey
+from superset.utils.core import get_example_database
 from tests.test_app import app  # isort:skip
 import superset.views.utils
 from superset import (
@@ -817,7 +818,9 @@ def test_comments_in_sqlatable_query(self):
         clean_query = "SELECT '/* val 1 */' as c1, '-- val 2' as c2 FROM tbl"
         commented_query = "/* comment 1 */" + clean_query + "-- comment 2"
         table = SqlaTable(
-            table_name="test_comments_in_sqlatable_query_table", sql=commented_query
+            table_name="test_comments_in_sqlatable_query_table",
+            sql=commented_query,
+            database=get_example_database(),
         )
         rendered_query = str(table.get_from_clause())
         self.assertEqual(clean_query, rendered_query)

tests/db_engine_specs/base_engine_spec_tests.py

@@ -14,19 +14,20 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from tests.test_app import app  # isort:skip
-
 import datetime
 from unittest import mock
 
 from superset.db_engine_specs import engines
 from superset.db_engine_specs.base import BaseEngineSpec, builtin_time_grains
 from superset.db_engine_specs.sqlite import SqliteEngineSpec
+from superset.sql_parse import ParsedQuery
 from superset.utils.core import get_example_database
 from tests.db_engine_specs.base_tests import TestDbEngineSpec
 
 from ..fixtures.pyodbcRow import Row
 
+from tests.test_app import app  # isort:skip
+
 
 class TestDbEngineSpecs(TestDbEngineSpec):
     def test_extract_limit_from_query(self, engine_spec_class=BaseEngineSpec):
@@ -239,3 +240,15 @@ def test_pyodbc_rows_to_tuples_passthrough(self):
         ]
         result = BaseEngineSpec.pyodbc_rows_to_tuples(data)
         self.assertListEqual(result, data)
+
+
+def test_is_readonly():
+    def is_readonly(sql: str) -> bool:
+        return BaseEngineSpec.is_readonly_query(ParsedQuery(sql))
+
+    assert not is_readonly("SHOW LOCKS test EXTENDED")
+    assert not is_readonly("SET hivevar:desc='Legislators'")
+    assert not is_readonly("UPDATE t1 SET col1 = NULL")
+    assert is_readonly("EXPLAIN SELECT 1")
+    assert is_readonly("SELECT 1")
+    assert is_readonly("WITH (SELECT 1) bla SELECT * from bla")

tests/db_engine_specs/hive_tests.py

@@ -23,7 +23,7 @@
 from tests.test_app import app
 from superset.db_engine_specs.hive import HiveEngineSpec
 from superset.exceptions import SupersetException
-from superset.sql_parse import Table
+from superset.sql_parse import Table, ParsedQuery
 
 
 def test_0_progress():
@@ -234,3 +234,16 @@ def test_get_create_table_stmt() -> None:
                 tblproperties ('skip.header.line.count'=:header_line_count)""",
         {"delim": ",", "location": "s3a://directory/table", "header_line_count": "101"},
     )
+
+
+def test_is_readonly():
+    def is_readonly(sql: str) -> bool:
+        return HiveEngineSpec.is_readonly_query(ParsedQuery(sql))
+
+    assert not is_readonly("UPDATE t1 SET col1 = NULL")
+    assert not is_readonly("INSERT OVERWRITE TABLE tabB SELECT a.Age FROM TableA")
+    assert is_readonly("SHOW LOCKS test EXTENDED")
+    assert is_readonly("SET hivevar:desc='Legislators'")
+    assert is_readonly("EXPLAIN SELECT 1")
+    assert is_readonly("SELECT 1")
+    assert is_readonly("WITH (SELECT 1) bla SELECT * from bla")

tests/db_engine_specs/presto_tests.py

@@ -22,6 +22,7 @@
 from sqlalchemy.sql import select
 
 from superset.db_engine_specs.presto import PrestoEngineSpec
+from superset.sql_parse import ParsedQuery
 from tests.db_engine_specs.base_tests import TestDbEngineSpec
 
 
@@ -514,3 +515,16 @@ def test_get_sqla_column_type(self):
 
         sqla_type = PrestoEngineSpec.get_sqla_column_type(None)
         assert sqla_type is None
+
+
+def test_is_readonly():
+    def is_readonly(sql: str) -> bool:
+        return PrestoEngineSpec.is_readonly_query(ParsedQuery(sql))
+
+    assert not is_readonly("SET hivevar:desc='Legislators'")
+    assert not is_readonly("UPDATE t1 SET col1 = NULL")
+    assert not is_readonly("INSERT OVERWRITE TABLE tabB SELECT a.Age FROM TableA")
+    assert is_readonly("SHOW LOCKS test EXTENDED")
+    assert is_readonly("EXPLAIN SELECT 1")
+    assert is_readonly("SELECT 1")
+    assert is_readonly("WITH (SELECT 1) bla SELECT * from bla")

tests/sql_parse_tests.py

@@ -376,14 +376,43 @@ def test_multistatement(self):
     def test_update_not_select(self):
         sql = ParsedQuery("UPDATE t1 SET col1 = NULL")
         self.assertEqual(False, sql.is_select())
-        self.assertEqual(False, sql.is_readonly())
+
+    def test_set(self):
+        sql = ParsedQuery(
+            """
+            -- comment
+            SET hivevar:desc='Legislators';
+        """
+        )
+
+        self.assertEqual(True, sql.is_set())
+        self.assertEqual(False, sql.is_select())
+
+        self.assertEqual(True, ParsedQuery("set hivevar:desc='bla'").is_set())
+        self.assertEqual(False, ParsedQuery("SELECT 1").is_set())
+
+    def test_show(self):
+        sql = ParsedQuery(
+            """
+            -- comment
+            SHOW LOCKS test EXTENDED;
+            -- comment
+        """
+        )
+
+        self.assertEqual(True, sql.is_show())
+        self.assertEqual(False, sql.is_select())
+
+        self.assertEqual(True, ParsedQuery("SHOW TABLES").is_show())
+        self.assertEqual(True, ParsedQuery("shOw TABLES").is_show())
+        self.assertEqual(True, ParsedQuery("show TABLES").is_show())
+        self.assertEqual(False, ParsedQuery("SELECT 1").is_show())
 
     def test_explain(self):
         sql = ParsedQuery("EXPLAIN SELECT 1")
 
         self.assertEqual(True, sql.is_explain())
         self.assertEqual(False, sql.is_select())
-        self.assertEqual(True, sql.is_readonly())
 
     def test_complex_extract_tables(self):
         query = """SELECT sum(m_examples) AS "sum__m_example"