feat: Implement MSSQL engine import validator for driver configuration and add corresponding tests

Author: mattiasthalen

sqlmesh/core/config/connection.py

@@ -1438,7 +1438,46 @@ class MSSQLConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["mssql"] = Field(alias="type", default="mssql")
 
-    _engine_import_validator = _get_engine_import_validator("pymssql", "mssql")
+    @model_validator(mode="before")
+    @classmethod
+    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
+
+            logger.exception("Failed to import the MSSQL engine library")
+
+            if driver == "pymssql":
+                extra_name = "mssql"
+            elif driver == "pyodbc":
+                extra_name = "mssql-odbc"
+            else:
+                extra_name = "mssql"
+
+            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."
+            )
+
+        return data
 
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:

tests/core/test_connection_config.py

@@ -4,6 +4,7 @@
 
 import pytest
 from _pytest.fixtures import FixtureRequest
+from unittest.mock import patch
 
 from sqlmesh.core.config.connection import (
     BigQueryConnectionConfig,
@@ -19,6 +20,7 @@
     SnowflakeConnectionConfig,
     TrinoAuthenticationMethod,
     AthenaConnectionConfig,
+    MSSQLConnectionConfig,
     _connection_config_validator,
     _get_engine_import_validator,
 )
@@ -1127,3 +1129,54 @@ class TestConfigC(PydanticModel):
         _engine_import_validator = _get_engine_import_validator("sqlmesh", "bigquery")
 
     TestConfigC()
+
+
+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
+        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
+        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
+        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
+    with patch("importlib.import_module") as mock_import:
+        mock_import.return_value = None
+        config = MSSQLConnectionConfig(host="localhost", driver="pyodbc")
+        assert config.driver == "pyodbc"