feat: Refactor MSSQL connection driver import handling

Author: mattiasthalen

sqlmesh/core/config/connection.py

@@ -1444,40 +1444,26 @@ def _mssql_engine_import_validator(cls, data: t.Any) -> t.Any:
         if not isinstance(data, dict):
             return data
 
-        check_import = str_to_bool(str(data.pop("check_import", True)))
-        if not check_import:
-            return data
-
         driver = data.get("driver", "pymssql")
 
-        try:
-            if driver == "pymssql":
-                importlib.import_module("pymssql")
-            elif driver == "pyodbc":
-                importlib.import_module("pyodbc")
-            else:
-                raise ValueError(f"Unsupported driver: {driver}")
-        except ImportError:
-            if debug_mode_enabled():
-                raise
+        # Define the mapping of driver to import module and extra name
+        driver_configs = {"pymssql": ("pymssql", "mssql"), "pyodbc": ("pyodbc", "mssql-odbc")}
 
-            logger.exception("Failed to import the MSSQL engine library")
+        if driver not in driver_configs:
+            raise ValueError(f"Unsupported driver: {driver}")
 
-            if driver == "pymssql":
-                extra_name = "mssql"
-            elif driver == "pyodbc":
-                extra_name = "mssql-odbc"
-            else:
-                extra_name = "mssql"
+        import_module, extra_name = driver_configs[driver]
 
-            raise ConfigError(
-                f"Failed to import the '{driver}' library for MSSQL connections. This may be due to a missing "
-                "or incompatible installation. Please ensure the required dependency is installed by "
-                f'running: `pip install "sqlmesh[{extra_name}]"`. For more details, check the logs '
-                "in the 'logs/' folder, or rerun the command with the '--debug' flag."
-            )
+        # Conditionally delegate to the existing _get_engine_import_validator
+        # Create a validator for the specific driver and call its inner function
+        validator_func = _get_engine_import_validator(import_module, driver, extra_name)
 
-        return data
+        # Extract the inner validate function from the decorated validator
+        # The validator_func has a __wrapped__ attribute that contains the original function
+        inner_validate = getattr(validator_func, "__wrapped__", validator_func)
+
+        # Call the inner validation function directly
+        return inner_validate(cls, data)
 
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:

tests/core/test_connection_config.py

@@ -1133,49 +1133,26 @@ class TestConfigC(PydanticModel):
 
 def test_mssql_engine_import_validator():
     """Test that MSSQL import validator respects driver configuration."""
-    with pytest.raises(
-        ConfigError,
-        match=re.escape(
-            "Failed to import the 'pyodbc' library for MSSQL connections. This may be due to a missing "
-            "or incompatible installation. Please ensure the required dependency is installed by "
-            'running: `pip install "sqlmesh[mssql-odbc]"`. For more details, check the logs '
-            "in the 'logs/' folder, or rerun the command with the '--debug' flag."
-        ),
-    ):
-        # Test PyODBC driver suggests mssql-odbc extra
+
+    # Test PyODBC driver suggests mssql-odbc extra when import fails
+    with pytest.raises(ConfigError, match=r"pip install \"sqlmesh\[mssql-odbc\]\""):
         with patch("importlib.import_module") as mock_import:
             mock_import.side_effect = ImportError("No module named 'pyodbc'")
             MSSQLConnectionConfig(host="localhost", driver="pyodbc")
 
-    with pytest.raises(
-        ConfigError,
-        match=re.escape(
-            "Failed to import the 'pymssql' library for MSSQL connections. This may be due to a missing "
-            "or incompatible installation. Please ensure the required dependency is installed by "
-            'running: `pip install "sqlmesh[mssql]"`. For more details, check the logs '
-            "in the 'logs/' folder, or rerun the command with the '--debug' flag."
-        ),
-    ):
-        # Test PyMSSQL driver suggests mssql extra
+    # Test PyMSSQL driver suggests mssql extra when import fails
+    with pytest.raises(ConfigError, match=r"pip install \"sqlmesh\[mssql\]\""):
         with patch("importlib.import_module") as mock_import:
             mock_import.side_effect = ImportError("No module named 'pymssql'")
             MSSQLConnectionConfig(host="localhost", driver="pymssql")
 
-    with pytest.raises(
-        ConfigError,
-        match=re.escape(
-            "Failed to import the 'pymssql' library for MSSQL connections. This may be due to a missing "
-            "or incompatible installation. Please ensure the required dependency is installed by "
-            'running: `pip install "sqlmesh[mssql]"`. For more details, check the logs '
-            "in the 'logs/' folder, or rerun the command with the '--debug' flag."
-        ),
-    ):
-        # Test default driver (pymssql) suggests mssql extra
+    # Test default driver (pymssql) suggests mssql extra when import fails
+    with pytest.raises(ConfigError, match=r"pip install \"sqlmesh\[mssql\]\""):
         with patch("importlib.import_module") as mock_import:
             mock_import.side_effect = ImportError("No module named 'pymssql'")
             MSSQLConnectionConfig(host="localhost")  # No driver specified
 
-    # Test successful import doesn't raise exception
+    # Test successful import works without error
     with patch("importlib.import_module") as mock_import:
         mock_import.return_value = None
         config = MSSQLConnectionConfig(host="localhost", driver="pyodbc")