feat(database): Database Filtering via custom configuration (#24580)

tests/unit_tests/databases/api_test.py

@@ -20,9 +20,11 @@
 import json
 from io import BytesIO
 from typing import Any
+from unittest.mock import Mock
 from uuid import UUID
 
 import pytest
+from flask import current_app
 from pytest_mock import MockFixture
 from sqlalchemy.orm.session import Session
 
@@ -495,3 +497,100 @@ def test_delete_ssh_tunnel_not_found(
 
         response_tunnel = DatabaseDAO.get_ssh_tunnel(2)
         assert response_tunnel is None
+
+
+def test_apply_dynamic_database_filter(
+    mocker: MockFixture,
+    app: Any,
+    session: Session,
+    client: Any,
+    full_api_access: None,
+) -> None:
+    """
+    Test that we can filter the list of databases.
+    First test the default behavior without a filter and then
+    defining a filter function and patching the config to get
+    the filtered results.
+    """
+    with app.app_context():
+        from superset.daos.database import DatabaseDAO
+        from superset.databases.api import DatabaseRestApi
+        from superset.databases.ssh_tunnel.models import SSHTunnel
+        from superset.models.core import Database
+
+        DatabaseRestApi.datamodel.session = session
+
+        # create table for databases
+        Database.metadata.create_all(session.get_bind())  # pylint: disable=no-member
+
+        # Create our First Database
+        database = Database(
+            database_name="first-database",
+            sqlalchemy_uri="gsheets://",
+            encrypted_extra=json.dumps(
+                {
+                    "metadata_params": {},
+                    "engine_params": {},
+                    "metadata_cache_timeout": {},
+                    "schemas_allowed_for_file_upload": [],
+                }
+            ),
+        )
+        session.add(database)
+        session.commit()
+
+        # Create our Second Database
+        database = Database(
+            database_name="second-database",
+            sqlalchemy_uri="gsheets://",
+            encrypted_extra=json.dumps(
+                {
+                    "metadata_params": {},
+                    "engine_params": {},
+                    "metadata_cache_timeout": {},
+                    "schemas_allowed_for_file_upload": [],
+                }
+            ),
+        )
+        session.add(database)
+        session.commit()
+
+        # mock the lookup so that we don't need to include the driver
+        mocker.patch("sqlalchemy.engine.URL.get_driver_name", return_value="gsheets")
+        mocker.patch("superset.utils.log.DBEventLogger.log")
+        mocker.patch(
+            "superset.databases.ssh_tunnel.commands.delete.is_feature_enabled",
+            return_value=False,
+        )
+
+        def _base_filter(query):
+            from superset.models.core import Database
+
+            return query.filter(Database.database_name.startswith("second"))
+
+        # Create a mock object
+        base_filter_mock = Mock(side_effect=_base_filter)
+
+        # Get our recently created Databases
+        response_databases = DatabaseDAO.find_all()
+        assert response_databases
+        expected_db_names = ["first-database", "second-database"]
+        actual_db_names = [db.database_name for db in response_databases]
+        assert actual_db_names == expected_db_names
+
+        # Ensure that the filter has not been called because it's not in our config
+        assert base_filter_mock.call_count == 0
+
+        original_config = current_app.config.copy()
+        original_config["EXTRA_DYNAMIC_QUERY_FILTERS"] = {"databases": base_filter_mock}
+
+        mocker.patch("superset.views.filters.current_app.config", new=original_config)
+        # Get filtered list
+        response_databases = DatabaseDAO.find_all()
+        assert response_databases
+        expected_db_names = ["second-database"]
+        actual_db_names = [db.database_name for db in response_databases]
+        assert actual_db_names == expected_db_names
+
+        # Ensure that the filter has been called once
+        assert base_filter_mock.call_count == 1