From 8c7e0f0197770d9b3a56451a38cc0eef5909d389 Mon Sep 17 00:00:00 2001 From: Daniel Gaspar Date: Tue, 18 Apr 2023 16:11:47 +0100 Subject: [PATCH 1/3] feat: add enforce URI query params with a specific for MySQL --- superset/db_engine_specs/base.py | 8 +++-- superset/db_engine_specs/mysql.py | 6 +++- tests/integration_tests/conftest.py | 2 +- tests/integration_tests/model_tests.py | 15 ++++++++ .../unit_tests/db_engine_specs/test_mysql.py | 34 +++++++++++++++++-- 5 files changed, 58 insertions(+), 7 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index ed58e8cb86b38..87a00af4d1a0f 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -357,6 +357,8 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods top_keywords: Set[str] = {"TOP"} # A set of disallowed connection query parameters disallow_uri_query_params: Set[str] = set() + # A Dict of query parameters that will always be used on every connection + enforce_uri_query_params: Dict[str, Any] = {} force_column_alias_quotes = False arraysize = 0 @@ -1089,11 +1091,11 @@ def adjust_engine_params( # pylint: disable=unused-argument ``supports_dynamic_schema`` set to true, so that Superset knows in which schema a given query is running in order to enforce permissions (see #23385 and #23401). - Currently, changing the catalog is not supported. The method acceps a catalog so - that when catalog support is added to Superse the interface remains the same. This + Currently, changing the catalog is not supported. The method accepts a catalog so + that when catalog support is added to Superset the interface remains the same. This is important because DB engine specs can be installed from 3rd party packages. """ - return uri, connect_args + return uri, {**connect_args, **cls.enforce_uri_query_params} @classmethod def patch(cls) -> None: diff --git a/superset/db_engine_specs/mysql.py b/superset/db_engine_specs/mysql.py index e5ff964f868da..07d2aea36286e 100644 --- a/superset/db_engine_specs/mysql.py +++ b/superset/db_engine_specs/mysql.py @@ -176,6 +176,7 @@ class MySQLEngineSpec(BaseEngineSpec, BasicParametersMixin): ), } disallow_uri_query_params = {"local_infile"} + enforce_uri_query_params = {"local_infile": 0} @classmethod def convert_dttm( @@ -198,10 +199,13 @@ def adjust_engine_params( catalog: Optional[str] = None, schema: Optional[str] = None, ) -> Tuple[URL, Dict[str, Any]]: + uri, new_connect_args = super( + MySQLEngineSpec, MySQLEngineSpec + ).adjust_engine_params(uri, connect_args, catalog, schema) if schema: uri = uri.set(database=parse.quote(schema, safe="")) - return uri, connect_args + return uri, new_connect_args @classmethod def get_schema_from_engine_params( diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index 0ea5bb5106b15..fdc015b94c942 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -121,7 +121,7 @@ def setup_sample_data() -> Any: # relying on `tests.integration_tests.test_app.app` leveraging an `app` fixture which is purposely # scoped to the function level to ensure tests remain idempotent. with app.app_context(): - setup_presto_if_needed() + # setup_presto_if_needed() from superset.cli.test import load_test_users_run diff --git a/tests/integration_tests/model_tests.py b/tests/integration_tests/model_tests.py index da6c5e6a3c254..35dbcc0a6bb3a 100644 --- a/tests/integration_tests/model_tests.py +++ b/tests/integration_tests/model_tests.py @@ -188,6 +188,21 @@ def test_impersonate_user_presto(self, mocked_create_engine): "password": "original_user_password", } + @unittest.skipUnless( + SupersetTestCase.is_module_installed("MySQLdb"), "mysqlclient not installed" + ) + @mock.patch("superset.models.core.create_engine") + def test_adjust_engine_params_mysql(self, mocked_create_engine): + model = Database( + database_name="test_database", + sqlalchemy_uri="mysql://user:password@localhost", + ) + model._get_sqla_engine() + call_args = mocked_create_engine.call_args + + assert str(call_args[0][0]) == "mysql://user:password@localhost" + assert call_args[1]["connect_args"]["local_infile"] == 0 + @mock.patch("superset.models.core.create_engine") def test_impersonate_user_trino(self, mocked_create_engine): principal_user = security_manager.find_user(username="gamma") diff --git a/tests/unit_tests/db_engine_specs/test_mysql.py b/tests/unit_tests/db_engine_specs/test_mysql.py index 091cdb3b46305..31e01ace58552 100644 --- a/tests/unit_tests/db_engine_specs/test_mysql.py +++ b/tests/unit_tests/db_engine_specs/test_mysql.py @@ -16,7 +16,7 @@ # under the License. from datetime import datetime -from typing import Any, Dict, Optional, Type +from typing import Any, Dict, Optional, Tuple, Type from unittest.mock import Mock, patch import pytest @@ -33,7 +33,7 @@ TINYINT, TINYTEXT, ) -from sqlalchemy.engine.url import make_url +from sqlalchemy.engine.url import make_url, URL from superset.utils.core import GenericDataType from tests.unit_tests.db_engine_specs.utils import ( @@ -119,6 +119,36 @@ def test_validate_database_uri(sqlalchemy_uri: str, error: bool) -> None: MySQLEngineSpec.validate_database_uri(url) +@pytest.mark.parametrize( + "sqlalchemy_uri,connect_args,returns", + [ + ("mysql://user:password@host/db1", {"local_infile": 1}, {"local_infile": 0}), + ("mysql://user:password@host/db1", {"local_infile": -1}, {"local_infile": 0}), + ("mysql://user:password@host/db1", {"local_infile": 0}, {"local_infile": 0}), + ( + "mysql://user:password@host/db1", + {"param1": "some_value"}, + {"local_infile": 0, "param1": "some_value"}, + ), + ( + "mysql://user:password@host/db1", + {"local_infile": 1, "param1": "some_value"}, + {"local_infile": 0, "param1": "some_value"}, + ), + ], +) +def test_adjust_engine_params( + sqlalchemy_uri: str, connect_args: Dict[str, Any], returns: Dict[str, Any] +) -> None: + from superset.db_engine_specs.mysql import MySQLEngineSpec + + url = make_url(sqlalchemy_uri) + returned_url, returned_connect_args = MySQLEngineSpec.adjust_engine_params( + url, connect_args + ) + assert returned_connect_args == returns + + @patch("sqlalchemy.engine.Engine.connect") def test_get_cancel_query_id(engine_mock: Mock) -> None: from superset.db_engine_specs.mysql import MySQLEngineSpec From 35ac9b3434a8ef6fcf75e4b047e78597f58a69f0 Mon Sep 17 00:00:00 2001 From: Daniel Gaspar Date: Tue, 18 Apr 2023 16:12:13 +0100 Subject: [PATCH 2/3] feat: add enforce URI query params with a specific for MySQL --- tests/integration_tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index fdc015b94c942..0ea5bb5106b15 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -121,7 +121,7 @@ def setup_sample_data() -> Any: # relying on `tests.integration_tests.test_app.app` leveraging an `app` fixture which is purposely # scoped to the function level to ensure tests remain idempotent. with app.app_context(): - # setup_presto_if_needed() + setup_presto_if_needed() from superset.cli.test import load_test_users_run From 7b61a48163c5aed6aa0882be855f0224b845b440 Mon Sep 17 00:00:00 2001 From: Daniel Gaspar Date: Tue, 18 Apr 2023 16:33:08 +0100 Subject: [PATCH 3/3] fix lint --- superset/db_engine_specs/base.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 87a00af4d1a0f..93df7c72168e9 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1092,8 +1092,9 @@ def adjust_engine_params( # pylint: disable=unused-argument given query is running in order to enforce permissions (see #23385 and #23401). Currently, changing the catalog is not supported. The method accepts a catalog so - that when catalog support is added to Superset the interface remains the same. This - is important because DB engine specs can be installed from 3rd party packages. + that when catalog support is added to Superset the interface remains the same. + This is important because DB engine specs can be installed from 3rd party + packages. """ return uri, {**connect_args, **cls.enforce_uri_query_params}