Adding DB_ENGINE_URI_VALIDATOR

superset/config.py

@@ -1206,6 +1206,17 @@ def CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC(  # pylint: disable=invalid-name
 DB_CONNECTION_MUTATOR = None
 
 
+# A callable that is invoked for every invocation of DB Engine Specs
+# which allows for custom validation of the engine URI.
+# See: superset.db_engine_specs.base.BaseEngineSpec.validate_database_uri
+# Example:
+#   def DB_ENGINE_URI_VALIDATOR(sqlalchemy_uri: URL):
+#       if not <some condition>:
+#           raise Exception("URI invalid")
+#
+DB_ENGINE_URI_VALIDATOR = None
+
+
 # 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
 # with information such as the username and worker node information

superset/db_engine_specs/base.py

@@ -1956,6 +1956,9 @@ def validate_database_uri(cls, sqlalchemy_uri: URL) -> None:
 
         :param sqlalchemy_uri:
         """
+        if db_engine_uri_validator := current_app.config["DB_ENGINE_URI_VALIDATOR"]:
+            db_engine_uri_validator(sqlalchemy_uri)
+
         if existing_disallowed := cls.disallow_uri_query_params.get(
             sqlalchemy_uri.get_driver_name(), set()
         ).intersection(sqlalchemy_uri.query):

tests/unit_tests/db_engine_specs/test_base.py

@@ -18,11 +18,13 @@
 
 from textwrap import dedent
 from typing import Any, Optional
+from unittest.mock import patch
 
 import pytest
 from pytest_mock import MockFixture
 from sqlalchemy import types
 from sqlalchemy.dialects import sqlite
+from sqlalchemy.engine.url import URL
 from sqlalchemy.sql import sqltypes
 
 from superset.superset_typing import ResultSetColumnType, SQLAColumnType
@@ -69,6 +71,23 @@ def test_parse_sql_multi_statement() -> None:
     ]
 
 
+@patch("superset.db_engine_specs.base.current_app")
+def test_validate_db_uri(current_app) -> None:
+    """
+    Ensures that the `validate_database_uri` method invokes the validator correctly
+    """
+
+    def mock_validate(sqlalchemy_uri: URL) -> None:
+        raise ValueError("Invalid URI")
+
+    current_app.config = {"DB_ENGINE_URI_VALIDATOR": mock_validate}
+
+    from superset.db_engine_specs.base import BaseEngineSpec
+
+    with pytest.raises(ValueError):
+        BaseEngineSpec.validate_database_uri(URL.create("sqlite"))
+
+
 @pytest.mark.parametrize(
     "original,expected",
     [