Fix: Improve error message for missing engine imports (#4447)

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
@@ -24,6 +25,7 @@
 )
 from sqlmesh.core.engine_adapter.shared import CatalogSupport
 from sqlmesh.core.engine_adapter import EngineAdapter
+from sqlmesh.utils import str_to_bool
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import (
     ValidationInfo,
@@ -49,6 +51,30 @@
 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:
+        check_import = (
+            str_to_bool(str(data.pop("check_import", True))) if isinstance(data, dict) else True
+        )
+        if not check_import:
+            return data
+        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 +454,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 +648,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 +901,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 +1045,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 +1176,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 +1237,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 +1290,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 +1342,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 +1399,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 +1517,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 +1669,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 +1775,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 +1843,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/engine_adapter/integration/config.yaml

@@ -17,6 +17,7 @@ gateways:
       catalog: datalake
       http_scheme: http
       retries: 20
+      check_import: false
     state_connection:
       type: duckdb
   inttest_trino_iceberg:
@@ -28,6 +29,7 @@ gateways:
       catalog: datalake_iceberg
       http_scheme: http
       retries: 20
+      check_import: false
     state_connection:
       type: duckdb
   inttest_trino_delta:
@@ -39,6 +41,7 @@ gateways:
       catalog: datalake_delta
       http_scheme: http
       retries: 20
+      check_import: false
     state_connection:
       type: duckdb
   inttest_trino_nessie:
@@ -50,13 +53,15 @@ gateways:
       catalog: datalake_nessie
       http_scheme: http
       retries: 20
+      check_import: false
     state_connection:
       type: duckdb
   inttest_spark:
     connection:
       type: spark
       config:
         spark.remote: sc://{{ env_var('DOCKER_HOSTNAME', 'localhost') }}
+      check_import: false
     state_connection:
       type: duckdb
   inttest_mssql:
@@ -65,6 +70,7 @@ gateways:
       host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}
       user: sa
       password: 1StrongPwd@@
+      check_import: false
   inttest_postgres:
     connection:
       type: postgres
@@ -73,6 +79,7 @@ gateways:
       database: postgres
       host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}
       port: 5432
+      check_import: false
   inttest_mysql:
     connection:
       type: mysql
@@ -81,13 +88,15 @@ gateways:
       password: mysql
       port: 3306
       charset: utf8
+      check_import: false
   inttest_clickhouse_standalone:
     connection:
       type: clickhouse
       host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}
       port: 8123
       username: clickhouse
       password: clickhouse
+      check_import: false
     state_connection:
       type: duckdb
   inttest_clickhouse_cluster:
@@ -98,6 +107,7 @@ gateways:
       username: clickhouse
       password: clickhouse
       cluster: cluster1
+      check_import: false
     state_connection:
       type: duckdb
   inttest_risingwave:
@@ -107,6 +117,7 @@ gateways:
       database: dev
       host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}
       port: 4566
+      check_import: false
 
 
   # Cloud databases
@@ -118,6 +129,7 @@ gateways:
       database: {{ env_var('SNOWFLAKE_DATABASE') }}
       user: {{ env_var('SNOWFLAKE_USER') }}
       password: {{ env_var('SNOWFLAKE_PASSWORD') }}
+      check_import: false
     state_connection:
       type: duckdb
 
@@ -128,6 +140,7 @@ gateways:
       server_hostname: {{ env_var('DATABRICKS_SERVER_HOSTNAME') }}
       http_path: {{ env_var('DATABRICKS_HTTP_PATH') }}
       access_token: {{ env_var('DATABRICKS_ACCESS_TOKEN') }}
+      check_import: false
 
   inttest_redshift:
     connection:
@@ -136,12 +149,14 @@ gateways:
       user: {{ env_var('REDSHIFT_USER') }}
       password: {{ env_var('REDSHIFT_PASSWORD') }}
       database: {{ env_var('REDSHIFT_DATABASE') }}
+      check_import: false
 
   inttest_bigquery:
     connection:
       type: bigquery
       method: service-account
       keyfile: {{ env_var('BIGQUERY_KEYFILE') }}
+      check_import: false
     state_connection:
       type: duckdb
 
@@ -155,6 +170,7 @@ gateways:
       connect_timeout: 30
       connection_pool_options:
         retries: 5
+      check_import: false
     state_connection:
       type: duckdb
 
@@ -166,6 +182,7 @@ gateways:
       region_name: {{ env_var("AWS_REGION") }}
       work_group: {{ env_var("ATHENA_WORK_GROUP", "primary") }}
       s3_warehouse_location: {{ env_var("ATHENA_S3_WAREHOUSE_LOCATION", "") }}
+      check_import: false
     state_connection:
       type: duckdb
 

tests/core/test_config.py

@@ -343,6 +343,7 @@ def test_load_yaml_config_env_var_gateway_override(tmp_path_factory):
         os.environ,
         {
             "SQLMESH__GATEWAYS__TESTING__STATE_CONNECTION__TYPE": "bigquery",
+            "SQLMESH__GATEWAYS__TESTING__STATE_CONNECTION__CHECK_IMPORT": "false",
             "SQLMESH__DEFAULT_GATEWAY": "testing",
         },
     ):
@@ -353,7 +354,7 @@ def test_load_yaml_config_env_var_gateway_override(tmp_path_factory):
             gateways={
                 "testing": GatewayConfig(
                     connection=MotherDuckConnectionConfig(database="blah"),
-                    state_connection=BigQueryConnectionConfig(),
+                    state_connection=BigQueryConnectionConfig(check_import=False),
                 ),
             },
             model_defaults=ModelDefaultsConfig(dialect="bigquery"),
@@ -373,6 +374,7 @@ def test_load_py_config_env_var_gateway_override(tmp_path_factory):
         os.environ,
         {
             "SQLMESH__GATEWAYS__DUCKDB_GATEWAY__STATE_CONNECTION__TYPE": "bigquery",
+            "SQLMESH__GATEWAYS__DUCKDB_GATEWAY__STATE_CONNECTION__CHECK_IMPORT": "false",
             "SQLMESH__DEFAULT_GATEWAY": "duckdb_gateway",
         },
     ):
@@ -384,7 +386,7 @@ def test_load_py_config_env_var_gateway_override(tmp_path_factory):
             gateways={  # type: ignore
                 "duckdb_gateway": GatewayConfig(
                     connection=DuckDBConnectionConfig(),
-                    state_connection=BigQueryConnectionConfig(),
+                    state_connection=BigQueryConnectionConfig(check_import=False),
                 ),
             },
             model_defaults=ModelDefaultsConfig(dialect=""),
@@ -823,6 +825,7 @@ def test_gcp_postgres_ip_and_scopes(tmp_path):
       gcp_postgres:
         connection:
           type: gcp_postgres
+          check_import: false
           instance_connection_string: something
           user: user
           password: password