Fix: Connection config parsing on the Airflow side (#1439)

Author: izeigerman

sqlmesh/core/config/__init__.py

@@ -12,6 +12,7 @@
     RedshiftConnectionConfig,
     SnowflakeConnectionConfig,
     SparkConnectionConfig,
+    parse_connection_config,
 )
 from sqlmesh.core.config.gateway import GatewayConfig
 from sqlmesh.core.config.loader import load_config_from_paths, load_config_from_yaml

sqlmesh/core/config/connection.py

@@ -807,12 +807,7 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
 }
 
 
-def _connection_config_validator(
-    cls: t.Type, v: ConnectionConfig | t.Dict[str, t.Any] | None
-) -> ConnectionConfig | None:
-    if v is None or isinstance(v, ConnectionConfig):
-        return v
-
+def parse_connection_config(v: t.Dict[str, t.Any]) -> ConnectionConfig:
     if "type" not in v:
         raise ConfigError("Missing connection type.")
 
@@ -823,6 +818,14 @@ def _connection_config_validator(
     return CONNECTION_CONFIG_TO_TYPE[connection_type](**v)
 
 
+def _connection_config_validator(
+    cls: t.Type, v: ConnectionConfig | t.Dict[str, t.Any] | None
+) -> ConnectionConfig | None:
+    if v is None or isinstance(v, ConnectionConfig):
+        return v
+    return parse_connection_config(v)
+
+
 connection_config_validator = field_validator(
     "connection",
     "state_connection",

sqlmesh/schedulers/airflow/util.py

@@ -6,7 +6,6 @@
 import typing as t
 from datetime import timedelta
 
-import pydantic
 from airflow import settings
 from airflow.api.common.experimental.delete_dag import delete_dag
 from airflow.exceptions import AirflowException, DagNotFound
@@ -17,7 +16,7 @@
 from sqlalchemy.orm import Session
 
 from sqlmesh.core import constants as c
-from sqlmesh.core.config import ConnectionConfig
+from sqlmesh.core.config import parse_connection_config
 from sqlmesh.core.engine_adapter import create_engine_adapter
 from sqlmesh.core.state_sync import EngineAdapterStateSync, StateSync
 from sqlmesh.schedulers.airflow import common
@@ -53,9 +52,7 @@ def scoped_state_sync() -> t.Iterator[StateSync]:
 
         logger.info("Using connection '%s' for state sync", connection.conn_id)
 
-        connection_config: ConnectionConfig = pydantic.parse_obj_as(
-            ConnectionConfig, connection_config_dict  # type: ignore
-        )
+        connection_config = parse_connection_config(connection_config_dict)
         engine_adapter = connection_config.create_engine_adapter()
     except AirflowException:
         logger.info("Using the Airflow database connection for state sync")