Fix: Serialization of the ConnectionConfig objects (#2214)

Author: izeigerman

sqlmesh/core/config/connection.py

@@ -24,6 +24,7 @@
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import (
+    PYDANTIC_MAJOR_VERSION,
     field_validator,
     model_validator,
     model_validator_v1_args,
@@ -42,9 +43,6 @@ class ConnectionConfig(abc.ABC, BaseConfig):
     concurrent_tasks: int
     register_comments: bool
 
-    # Pydantic V2 does not include type field in dict unless also in base class
-    type_: str = Field(alias="type", default="none")
-
     @property
     @abc.abstractmethod
     def _connection_kwargs_keys(self) -> t.Set[str]:
@@ -1276,3 +1274,16 @@ def _connection_config_validator(
     mode="before",
     check_fields=False,
 )(_connection_config_validator)
+
+
+if t.TYPE_CHECKING:
+    # TypeAlias hasn't been introduced until Python 3.10 which means that we can't use it
+    # outside the TYPE_CHECKING guard.
+    SerializableConnectionConfig: t.TypeAlias = ConnectionConfig  # type: ignore
+elif PYDANTIC_MAJOR_VERSION >= 2:
+    import pydantic
+
+    # Workaround for https://docs.pydantic.dev/latest/concepts/serialization/#serializing-with-duck-typing
+    SerializableConnectionConfig = pydantic.SerializeAsAny[ConnectionConfig]  # type: ignore
+else:
+    SerializableConnectionConfig = ConnectionConfig  # type: ignore

sqlmesh/core/config/gateway.py

@@ -4,7 +4,10 @@
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.config.base import BaseConfig
-from sqlmesh.core.config.connection import ConnectionConfig, connection_config_validator
+from sqlmesh.core.config.connection import (
+    SerializableConnectionConfig,
+    connection_config_validator,
+)
 from sqlmesh.core.config.scheduler import SchedulerConfig
 
 
@@ -22,9 +25,9 @@ class GatewayConfig(BaseConfig):
             then no schema name is used and therefore the default schema defined for the connection will be used
     """
 
-    connection: t.Optional[ConnectionConfig] = None
-    state_connection: t.Optional[ConnectionConfig] = None
-    test_connection: t.Optional[ConnectionConfig] = None
+    connection: t.Optional[SerializableConnectionConfig] = None
+    state_connection: t.Optional[SerializableConnectionConfig] = None
+    test_connection: t.Optional[SerializableConnectionConfig] = None
     scheduler: t.Optional[SchedulerConfig] = None
     state_schema: t.Optional[str] = c.SQLMESH
 

sqlmesh/core/config/root.py

@@ -17,6 +17,7 @@
 from sqlmesh.core.config.connection import (
     ConnectionConfig,
     DuckDBConnectionConfig,
+    SerializableConnectionConfig,
     connection_config_validator,
 )
 from sqlmesh.core.config.feature_flag import FeatureFlag
@@ -71,8 +72,8 @@ class Config(BaseConfig):
     """
 
     gateways: t.Dict[str, GatewayConfig] = {"": GatewayConfig()}
-    default_connection: ConnectionConfig = DuckDBConnectionConfig()
-    default_test_connection_: t.Optional[ConnectionConfig] = Field(
+    default_connection: SerializableConnectionConfig = DuckDBConnectionConfig()
+    default_test_connection_: t.Optional[SerializableConnectionConfig] = Field(
         default=None, alias="default_test_connection"
     )
     default_scheduler: SchedulerConfig = BuiltInSchedulerConfig()

tests/core/test_config.py

@@ -437,3 +437,27 @@ class DerivedConfig(Config):
         },
         model_defaults=ModelDefaultsConfig(dialect=""),
     )
+
+
+def test_connection_config_serialization():
+    config = Config(
+        default_connection=DuckDBConnectionConfig(database="my_db"),
+        default_test_connection=DuckDBConnectionConfig(database="my_test_db"),
+    )
+    serialized = config.dict()
+    assert serialized["default_connection"] == {
+        "concurrent_tasks": 1,
+        "register_comments": True,
+        "type": "duckdb",
+        "extensions": [],
+        "connector_config": {},
+        "database": "my_db",
+    }
+    assert serialized["default_test_connection"] == {
+        "concurrent_tasks": 1,
+        "register_comments": True,
+        "type": "duckdb",
+        "extensions": [],
+        "connector_config": {},
+        "database": "my_test_db",
+    }