fix: adds the ability to disallow SQL functions per engine (#28639)

Author: dpgaspar

superset-frontend/cypress-base/cypress.config.ts

@@ -29,7 +29,7 @@ export default defineConfig({
   videoUploadOnPasses: false,
   viewportWidth: 1280,
   viewportHeight: 1024,
-  projectId: 'ukwxzo',
+  projectId: 'ud5x2f',
   retries: {
     runMode: 2,
     openMode: 0,

superset-frontend/src/dashboard/components/Header/HeaderActionsDropdown/index.jsx

@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React from 'react';
+import React, { PureComponent } from 'react';
 import PropTypes from 'prop-types';
 import { isEmpty } from 'lodash';
 import { connect } from 'react-redux';

superset/config.py

@@ -1203,6 +1203,15 @@ def CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC(  # pylint: disable=invalid-name
 # as such `create_engine(url, **params)`
 DB_CONNECTION_MUTATOR = None
 
+# A set of disallowed SQL functions per engine. This is used to restrict the use of
+# unsafe SQL functions in SQL Lab and Charts. The keys of the dictionary are the engine
+# names, and the values are sets of disallowed functions.
+DISALLOWED_SQL_FUNCTIONS: dict[str, set[str]] = {
+    "postgresql": {"version", "query_to_xml", "inet_server_addr", "inet_client_addr"},
+    "clickhouse": {"url"},
+    "mysql": {"version"},
+}
+
 
 # A function that intercepts the SQL to be executed and can alter it.
 # The use case is can be around adding some sort of comment header

superset/db_engine_specs/base.py

@@ -59,6 +59,7 @@
 from superset.constants import TimeGrain as TimeGrainConstants
 from superset.databases.utils import make_url_safe
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import DisallowedSQLFunction
 from superset.sql_parse import ParsedQuery, Table
 from superset.superset_typing import ResultSetColumnType, SQLAColumnType
 from superset.utils import core as utils
@@ -1584,6 +1585,11 @@ def execute(  # pylint: disable=unused-argument
         """
         if not cls.allows_sql_comments:
             query = sql_parse.strip_comments_from_sql(query, engine=cls.engine)
+        disallowed_functions = current_app.config["DISALLOWED_SQL_FUNCTIONS"].get(
+            cls.engine, set()
+        )
+        if sql_parse.check_sql_functions_exist(query, disallowed_functions, cls.engine):
+            raise DisallowedSQLFunction(disallowed_functions)
 
         if cls.arraysize:
             cursor.arraysize = cls.arraysize

superset/db_engine_specs/trino.py

@@ -23,7 +23,7 @@
 from typing import Any, TYPE_CHECKING
 
 import simplejson as json
-from flask import current_app
+from flask import current_app, Flask
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.engine.url import URL
 from sqlalchemy.exc import NoSuchTableError
@@ -206,19 +206,27 @@ def execute_with_cursor(cls, cursor: Cursor, sql: str, query: Query) -> None:
         execute_result: dict[str, Any] = {}
         execute_event = threading.Event()
 
-        def _execute(results: dict[str, Any], event: threading.Event) -> None:
+        def _execute(
+            results: dict[str, Any], event: threading.Event, app: Flask
+        ) -> None:
             logger.debug("Query %d: Running query: %s", query_id, sql)
 
             try:
-                cls.execute(cursor, sql)
+                with app.app_context():
+                    cls.execute(cursor, sql)
             except Exception as ex:  # pylint: disable=broad-except
                 results["error"] = ex
             finally:
                 event.set()
 
         execute_thread = threading.Thread(
             target=_execute,
-            args=(execute_result, execute_event),
+            args=(
+                execute_result,
+                execute_event,
+                # pylint: disable=protected-access
+                current_app._get_current_object(),
+            ),
         )
         execute_thread.start()
 

superset/exceptions.py

@@ -295,3 +295,18 @@ def __init__(self, exc: ValidationError, payload: dict[str, Any]):
             extra={"messages": exc.messages, "payload": payload},
         )
         super().__init__(error)
+
+
+class DisallowedSQLFunction(SupersetErrorException):
+    """
+    Disallowed function found on SQL statement
+    """
+
+    def __init__(self, functions: set[str]):
+        super().__init__(
+            SupersetError(
+                message=f"SQL statement contains disallowed function(s): {functions}",
+                error_type=SupersetErrorType.SYNTAX_ERROR,
+                level=ErrorLevel.ERROR,
+            )
+        )

superset/sql_parse.py

@@ -36,6 +36,7 @@
 from sqlparse import keywords
 from sqlparse.lexer import Lexer
 from sqlparse.sql import (
+    Function,
     Identifier,
     IdentifierList,
     Parenthesis,
@@ -219,6 +220,19 @@ def get_cte_remainder_query(sql: str) -> tuple[str | None, str]:
     return cte, remainder
 
 
+def check_sql_functions_exist(
+    sql: str, function_list: set[str], engine: str | None = None
+) -> bool:
+    """
+    Check if the SQL statement contains any of the specified functions.
+
+    :param sql: The SQL statement
+    :param function_list: The list of functions to search for
+    :param engine: The engine to use for parsing the SQL statement
+    """
+    return ParsedQuery(sql, engine=engine).check_functions_exist(function_list)
+
+
 def strip_comments_from_sql(statement: str, engine: str | None = None) -> str:
     """
     Strips comments from a SQL statement, does a simple test first
@@ -288,6 +302,34 @@ def tables(self) -> set[Table]:
             self._tables = self._extract_tables_from_sql()
         return self._tables
 
+    def _check_functions_exist_in_token(
+        self, token: Token, functions: set[str]
+    ) -> bool:
+        if (
+            isinstance(token, Function)
+            and token.get_name() is not None
+            and token.get_name().lower() in functions
+        ):
+            return True
+        if hasattr(token, "tokens"):
+            for inner_token in token.tokens:
+                if self._check_functions_exist_in_token(inner_token, functions):
+                    return True
+        return False
+
+    def check_functions_exist(self, functions: set[str]) -> bool:
+        """
+        Check if the SQL statement contains any of the specified functions.
+
+        :param functions: A set of functions to search for
+        :return: True if the statement contains any of the specified functions
+        """
+        for statement in self._parsed:
+            for token in statement.tokens:
+                if self._check_functions_exist_in_token(token, functions):
+                    return True
+        return False
+
     def _extract_tables_from_sql(self) -> set[Table]:
         """
         Extract all table references in a query.

tests/unit_tests/db_engine_specs/test_trino.py

@@ -396,7 +396,7 @@ def test_handle_cursor_early_cancel(
         assert cancel_query_mock.call_args is None
 
 
-def test_execute_with_cursor_in_parallel(mocker: MockerFixture):
+def test_execute_with_cursor_in_parallel(app, mocker: MockerFixture):
     """Test that `execute_with_cursor` fetches query ID from the cursor"""
     from superset.db_engine_specs.trino import TrinoEngineSpec
 
@@ -411,16 +411,20 @@ def _mock_execute(*args, **kwargs):
         mock_cursor.query_id = query_id
 
     mock_cursor.execute.side_effect = _mock_execute
+    with patch.dict(
+        "superset.config.DISALLOWED_SQL_FUNCTIONS",
+        {},
+        clear=True,
+    ):
+        TrinoEngineSpec.execute_with_cursor(
+            cursor=mock_cursor,
+            sql="SELECT 1 FROM foo",
+            query=mock_query,
+        )
 
-    TrinoEngineSpec.execute_with_cursor(
-        cursor=mock_cursor,
-        sql="SELECT 1 FROM foo",
-        query=mock_query,
-    )
-
-    mock_query.set_extra_json_key.assert_called_once_with(
-        key=QUERY_CANCEL_KEY, value=query_id
-    )
+        mock_query.set_extra_json_key.assert_called_once_with(
+            key=QUERY_CANCEL_KEY, value=query_id
+        )
 
 
 def test_get_columns(mocker: MockerFixture):

tests/unit_tests/sql_parse_tests.py

@@ -32,6 +32,7 @@
 )
 from superset.sql_parse import (
     add_table_name,
+    check_sql_functions_exist,
     extract_table_references,
     extract_tables_from_jinja_sql,
     get_rls_for_table,
@@ -1292,6 +1293,31 @@ def test_strip_comments_from_sql() -> None:
     )
 
 
+def test_check_sql_functions_exist() -> None:
+    """
+    Test that comments are stripped out correctly.
+    """
+    assert not (
+        check_sql_functions_exist("select a, b from version", {"version"}, "postgresql")
+    )
+
+    assert check_sql_functions_exist("select version()", {"version"}, "postgresql")
+
+    assert check_sql_functions_exist(
+        "select version from version()", {"version"}, "postgresql"
+    )
+
+    assert check_sql_functions_exist(
+        "select 1, a.version from (select version from version()) as a",
+        {"version"},
+        "postgresql",
+    )
+
+    assert check_sql_functions_exist(
+        "select 1, a.version from (select version()) as a", {"version"}, "postgresql"
+    )
+
+
 def test_sanitize_clause_valid():
     # regular clauses
     assert sanitize_clause("col = 1") == "col = 1"

tests/unit_tests/utils/csv_tests.py

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import sys
+
 import pandas as pd
 import pyarrow as pa
 import pytest
@@ -57,6 +59,7 @@ def test_escape_value():
     assert result == "' =10+2"
 
 
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10 or later")
 def test_df_to_escaped_csv():
     df = pd.DataFrame(
         data={
@@ -87,7 +90,7 @@ def test_df_to_escaped_csv():
         ["col_a"],
         ["'=func()"],
         ["-10"],
-        ["\"'=cmd\\|' /C calc'!A0\""],
+        [r"'=cmd\\|' /C calc'!A0"],
         ['"\'""""=b"'],
         ["' =a"],
         ["\x00"],