Fix: Improve error message for missing engine imports

Author: izeigerman

sqlmesh/core/config/connection.py

@@ -4,6 +4,7 @@
 import base64
 import logging
 import os
+import importlib
 import pathlib
 import re
 import typing as t
@@ -49,6 +50,25 @@
 MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)")
 
 
+def _get_engine_import_validator(
+    import_name: str, engine_type: str, extra_name: t.Optional[str] = None
+) -> t.Callable:
+    extra_name = extra_name or engine_type
+
+    @model_validator(mode="before")
+    def validate(cls: t.Any, data: t.Any) -> t.Any:
+        try:
+            importlib.import_module(import_name)
+        except ImportError:
+            raise ConfigError(
+                f"Failed to import the '{engine_type}' engine library. Please run `pip install \"sqlmesh[{extra_name}]\"`."
+            )
+
+        return data
+
+    return validate
+
+
 class ConnectionConfig(abc.ABC, BaseConfig):
     type_: str
     concurrent_tasks: int
@@ -428,6 +448,7 @@ class SnowflakeConnectionConfig(ConnectionConfig):
     type_: t.Literal["snowflake"] = Field(alias="type", default="snowflake")
 
     _concurrent_tasks_validator = concurrent_tasks_validator
+    _engine_import_validator = _get_engine_import_validator("snowflake", "snowflake")
 
     @model_validator(mode="before")
     def _validate_authenticator(cls, data: t.Any) -> t.Any:
@@ -621,6 +642,7 @@ class DatabricksConnectionConfig(ConnectionConfig):
 
     _concurrent_tasks_validator = concurrent_tasks_validator
     _http_headers_validator = http_headers_validator
+    _engine_import_validator = _get_engine_import_validator("databricks", "databricks")
 
     @model_validator(mode="before")
     def _databricks_connect_validator(cls, data: t.Any) -> t.Any:
@@ -873,6 +895,8 @@ class BigQueryConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["bigquery"] = Field(alias="type", default="bigquery")
 
+    _engine_import_validator = _get_engine_import_validator("google.cloud.bigquery", "bigquery")
+
     @field_validator("execution_project")
     def validate_execution_project(
         cls,
@@ -1015,6 +1039,10 @@ class GCPPostgresConnectionConfig(ConnectionConfig):
     register_comments: bool = True
     pre_ping: bool = True
 
+    _engine_import_validator = _get_engine_import_validator(
+        "google.cloud.sql", "gcp_postgres", "gcppostgres"
+    )
+
     @model_validator(mode="before")
     def _validate_auth_method(cls, data: t.Any) -> t.Any:
         if not isinstance(data, dict):
@@ -1142,6 +1170,8 @@ class RedshiftConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["redshift"] = Field(alias="type", default="redshift")
 
+    _engine_import_validator = _get_engine_import_validator("redshift_connector", "redshift")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1201,6 +1231,8 @@ class PostgresConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["postgres"] = Field(alias="type", default="postgres")
 
+    _engine_import_validator = _get_engine_import_validator("psycopg2", "postgres")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1252,6 +1284,8 @@ class MySQLConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["mysql"] = Field(alias="type", default="mysql")
 
+    _engine_import_validator = _get_engine_import_validator("pymysql", "mysql")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         connection_keys = {
@@ -1302,6 +1336,8 @@ class MSSQLConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["mssql"] = Field(alias="type", default="mssql")
 
+    _engine_import_validator = _get_engine_import_validator("pymssql", "mssql")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1357,6 +1393,8 @@ class SparkConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["spark"] = Field(alias="type", default="spark")
 
+    _engine_import_validator = _get_engine_import_validator("pyspark", "spark")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {
@@ -1473,6 +1511,8 @@ class TrinoConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["trino"] = Field(alias="type", default="trino")
 
+    _engine_import_validator = _get_engine_import_validator("trino", "trino")
+
     @field_validator("schema_location_mapping", mode="before")
     @classmethod
     def _validate_regex_keys(
@@ -1623,6 +1663,8 @@ class ClickhouseConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["clickhouse"] = Field(alias="type", default="clickhouse")
 
+    _engine_import_validator = _get_engine_import_validator("clickhouse_connect", "clickhouse")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         kwargs = {
@@ -1727,6 +1769,8 @@ class AthenaConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["athena"] = Field(alias="type", default="athena")
 
+    _engine_import_validator = _get_engine_import_validator("pyathena", "athena")
+
     @model_validator(mode="after")
     def _root_validator(self) -> Self:
         work_group = self.work_group
@@ -1793,6 +1837,8 @@ class RisingwaveConnectionConfig(ConnectionConfig):
 
     type_: t.Literal["risingwave"] = Field(alias="type", default="risingwave")
 
+    _engine_import_validator = _get_engine_import_validator("psycopg2", "risingwave")
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return {

tests/core/test_connection_config.py

@@ -20,8 +20,10 @@
     TrinoAuthenticationMethod,
     AthenaConnectionConfig,
     _connection_config_validator,
+    _get_engine_import_validator,
 )
 from sqlmesh.utils.errors import ConfigError
+from sqlmesh.utils.pydantic import PydanticModel
 
 
 @pytest.fixture
@@ -994,3 +996,36 @@ def test_databricks(make_config):
             server_hostname="dbc-test.cloud.databricks.com",
             auth_type="databricks-oauth",
         )
+
+
+def test_engine_import_validator():
+    with pytest.raises(
+        ConfigError,
+        match=re.escape(
+            """Failed to import the 'bigquery' engine library. Please run `pip install "sqlmesh[bigquery]"`."""
+        ),
+    ):
+
+        class TestConfigA(PydanticModel):
+            _engine_import_validator = _get_engine_import_validator("missing", "bigquery")
+
+        TestConfigA()
+
+    with pytest.raises(
+        ConfigError,
+        match=re.escape(
+            """Failed to import the 'bigquery' engine library. Please run `pip install "sqlmesh[bigquery_extra]"`."""
+        ),
+    ):
+
+        class TestConfigB(PydanticModel):
+            _engine_import_validator = _get_engine_import_validator(
+                "missing", "bigquery", "bigquery_extra"
+            )
+
+        TestConfigB()
+
+    class TestConfigC(PydanticModel):
+        _engine_import_validator = _get_engine_import_validator("sqlmesh", "bigquery")
+
+    TestConfigC()