Repository: TobikoData/sqlmesh
Pull request #2214 (2024-03-04 20:38:34)

Fix: Serialization of the ConnectionConfig objects
The serialization of abstract base classes has been intentionally broken in Pydantic V2. See https://docs.pydantic.dev/latest/concepts/serialization/#serializing-with-duck-typing for more details.


Patch:
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 861c4e29..5819fef0 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -24,6 +24,7 @@ from sqlmesh.core.config.common import (
 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 @@ connection_config_validator = field_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
diff --git a/sqlmesh/core/config/gateway.py b/sqlmesh/core/config/gateway.py
index 11ab9507..e2ffaff2 100644
--- a/sqlmesh/core/config/gateway.py
+++ b/sqlmesh/core/config/gateway.py
@@ -4,7 +4,10 @@ import typing as t
 
 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
 
diff --git a/sqlmesh/core/config/root.py b/sqlmesh/core/config/root.py
index 922877bf..98cc7e71 100644
--- a/sqlmesh/core/config/root.py
+++ b/sqlmesh/core/config/root.py
@@ -17,6 +17,7 @@ from sqlmesh.core.config.base import BaseConfig, UpdateStrategy
 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()


Test patch:
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index 60494936..23576db4 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -437,3 +437,27 @@ def test_load_alternative_config_type(yaml_config_path: Path, python_config_path
         },
         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",
+    }