feat: add denylist for db engines (#21486)

superset/config.py

@@ -38,6 +38,7 @@ from typing import (
     List,
     Literal,
     Optional,
+    Set,
     Type,
     TYPE_CHECKING,
     Union,
@@ -1107,6 +1108,12 @@ def EMAIL_HEADER_MUTATOR(  # pylint: disable=invalid-name,unused-argument
 # in security manager
 EXCLUDE_USERS_FROM_LISTS: Optional[List[str]] = None
 
+# For database connections, this dictionary will remove engines from the available
+# list/dropdown if you do not want these dbs to show as available.
+# The available list is generated by driver installed, and some engines have multiple
+# drivers.
+# e.g., DBS_AVAILABLE_DENYLIST: Dict[str, Set[str]] = {"databricks": ("pyhive", "pyodbc")}
+DBS_AVAILABLE_DENYLIST: Dict[str, Set[str]] = {}
 
 # This auth provider is used by background (offline) tasks that need to access
 # protected resources. Can be overridden by end users in order to support

superset/db_engine_specs/__init__.py

@@ -41,6 +41,7 @@ from pkg_resources import iter_entry_points
 from sqlalchemy.engine.default import DefaultDialect
 from sqlalchemy.engine.url import URL
 
+from superset import app
 from superset.db_engine_specs.base import BaseEngineSpec
 
 logger = logging.getLogger(__name__)
@@ -72,7 +73,6 @@ def load_engine_specs() -> List[Type[BaseEngineSpec]]:
             for attr in module.__dict__
             if is_engine_spec(getattr(module, attr))
         )
-
     # load additional engines from external modules
     for ep in iter_entry_points("superset.db_engine_specs"):
         try:
@@ -170,6 +170,17 @@ def get_available_engine_specs() -> Dict[Type[BaseEngineSpec], Set[str]]:
     for engine_spec in load_engine_specs():
         driver = drivers[engine_spec.engine]
 
+        # do not add denied db engine specs to available list
+        dbs_denylist = app.config["DBS_AVAILABLE_DENYLIST"]
+        dbs_denylist_engines = dbs_denylist.keys()
+
+        if (
+            engine_spec.engine in dbs_denylist_engines
+            and hasattr(engine_spec, "default_driver")
+            and engine_spec.default_driver in dbs_denylist[engine_spec.engine]
+        ):
+            continue
+
         # lookup driver by engine aliases.
         if not driver and engine_spec.engine_aliases:
             for alias in engine_spec.engine_aliases:

tests/unit_tests/conftest.py

@@ -21,6 +21,7 @@ import os
 from typing import Any, Callable, Iterator
 
 import pytest
+from _pytest.fixtures import SubRequest
 from pytest_mock import MockFixture
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
@@ -68,7 +69,7 @@ def session(get_session) -> Iterator[Session]:
 
 
 @pytest.fixture(scope="module")
-def app() -> Iterator[SupersetApp]:
+def app(request: SubRequest) -> Iterator[SupersetApp]:
     """
     A fixture that generates a Superset app.
     """
@@ -82,6 +83,11 @@ def app() -> Iterator[SupersetApp]:
     app.config["PREVENT_UNSAFE_DB_CONNECTIONS"] = False
     app.config["TESTING"] = True
 
+    # loop over extra configs passed in by tests
+    if request and hasattr(request, "param"):
+        for key, val in request.param.items():
+            app.config[key] = val
+
     # ``superset.extensions.appbuilder`` is a singleton, and won't rebuild the
     # routes when this fixture is called multiple times; we need to clear the
     # registered views to ensure the initialization can happen more than once.

tests/unit_tests/db_engine_specs/test_init.py

@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+import pytest
+from pkg_resources import EntryPoint
+from pytest_mock import MockFixture
+
+from superset.db_engine_specs import get_available_engine_specs
+
+
+def test_get_available_engine_specs(mocker: MockFixture) -> None:
+    """
+    get_available_engine_specs should return all engine specs
+    """
+    from superset.db_engine_specs.databricks import (
+        DatabricksHiveEngineSpec,
+        DatabricksNativeEngineSpec,
+        DatabricksODBCEngineSpec,
+    )
+
+    mocker.patch(
+        "superset.db_engine_specs.load_engine_specs",
+        return_value=iter(
+            [
+                DatabricksHiveEngineSpec,
+                DatabricksNativeEngineSpec,
+                DatabricksODBCEngineSpec,
+            ]
+        ),
+    )
+
+    assert list(get_available_engine_specs().keys()) == [
+        DatabricksHiveEngineSpec,
+        DatabricksNativeEngineSpec,
+        DatabricksODBCEngineSpec,
+    ]
+
+
+@pytest.mark.parametrize(
+    "app",
+    [{"DBS_AVAILABLE_DENYLIST": {"databricks": {"pyhive", "pyodbc"}}}],
+    indirect=True,
+)
+def test_get_available_engine_specs_with_denylist(mocker: MockFixture) -> None:
+    """
+    The denylist removes items from the db engine spec list
+    """
+    from superset.db_engine_specs.databricks import (
+        DatabricksHiveEngineSpec,
+        DatabricksNativeEngineSpec,
+        DatabricksODBCEngineSpec,
+    )
+
+    mocker.patch(
+        "superset.db_engine_specs.load_engine_specs",
+        return_value=iter(
+            [
+                DatabricksHiveEngineSpec,
+                DatabricksNativeEngineSpec,
+                DatabricksODBCEngineSpec,
+            ]
+        ),
+    )
+    available = get_available_engine_specs()
+    assert list(available.keys()) == [DatabricksNativeEngineSpec]