Repository: TobikoData/sqlmesh
Pull request #3085 (2024-09-04 14:29:44)

fix!: properly nested update gateway configs
Prior to this change nested updates to Gateways would only work with YAML config files. This update makes it work with Python config too. 

`test_load_yaml_config_env_var_gateway_override` passed without any changes (confirming expected behavior) while `test_load_py_config_env_var_gateway_override` fails without the change. 


Patch:
diff --git a/sqlmesh/core/config/base.py b/sqlmesh/core/config/base.py
index 78d1ed70..0da36e47 100644
--- a/sqlmesh/core/config/base.py
+++ b/sqlmesh/core/config/base.py
@@ -36,6 +36,15 @@ def update_field(
 
         The updated field
     """
+
+    def _update_pydantic_config(old: BaseConfig, new: BaseConfig) -> PydanticModel:
+        if type(new) != type(old):
+            raise ConfigError(
+                "NESTED_UPDATE behavior requires both values to have the same type. "
+                f"{type(old)} and {type(new)} were given instead."
+            )
+        return old.update_with(new)
+
     if not old:
         return new
 
@@ -78,18 +87,20 @@ def update_field(
 
         return combined
     if update_strategy == UpdateStrategy.NESTED_UPDATE:
-        if not isinstance(old, BaseConfig):
+        if not isinstance(old, BaseConfig) and not isinstance(old, dict):
             raise ConfigError(
-                f"NESTED_UPDATE behavior requires a config object. {type(old)} was given instead."
+                f"NESTED_UPDATE behavior requires a config object and a dict of config objects as values. {type(old)} was given instead."
             )
 
-        if type(new) != type(old):
-            raise ConfigError(
-                "NESTED_UPDATE behavior requires both values to have the same type. "
-                f"{type(old)} and {type(new)} were given instead."
-            )
+        if isinstance(old, dict):
+            for k, pydantic_model in new.items():
+                if k in old:
+                    old[k] = _update_pydantic_config(old[k], pydantic_model)
+                else:
+                    old[k] = pydantic_model
 
-        return old.update_with(new)
+            return old
+        return _update_pydantic_config(old, new)
 
     raise ConfigError(f"Unknown update strategy {update_strategy}.")
 
diff --git a/sqlmesh/core/config/root.py b/sqlmesh/core/config/root.py
index a2e18e84..565ccc00 100644
--- a/sqlmesh/core/config/root.py
+++ b/sqlmesh/core/config/root.py
@@ -118,7 +118,7 @@ class Config(BaseConfig):
     disable_anonymized_analytics: bool = False
 
     _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
-        "gateways": UpdateStrategy.KEY_UPDATE,
+        "gateways": UpdateStrategy.NESTED_UPDATE,
         "notification_targets": UpdateStrategy.EXTEND,
         "ignore_patterns": UpdateStrategy.EXTEND,
         "users": UpdateStrategy.EXTEND,


Test patch:
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index 9e831731..b1577af4 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -12,6 +12,8 @@ from sqlmesh.core.config import (
     DuckDBConnectionConfig,
     GatewayConfig,
     ModelDefaultsConfig,
+    BigQueryConnectionConfig,
+    MotherDuckConnectionConfig,
 )
 from sqlmesh.core.config.connection import DuckDBAttachOptions
 from sqlmesh.core.config.feature_flag import DbtFeatureFlag, FeatureFlag
@@ -307,6 +309,73 @@ def test_load_config_from_env_invalid_variable_name():
             load_config_from_env()
 
 
+def test_load_yaml_config_env_var_gateway_override(tmp_path_factory):
+    config_path = tmp_path_factory.mktemp("yaml_config") / "config.yaml"
+    with open(config_path, "w", encoding="utf-8") as fd:
+        fd.write(
+            """
+gateways:
+    testing:
+        connection:
+            type: motherduck
+            database: blah
+model_defaults:
+    dialect: bigquery
+        """
+        )
+    with mock.patch.dict(
+        os.environ,
+        {
+            "SQLMESH__GATEWAYS__TESTING__STATE_CONNECTION__TYPE": "bigquery",
+            "SQLMESH__DEFAULT_GATEWAY": "testing",
+        },
+    ):
+        assert load_config_from_paths(
+            Config,
+            project_paths=[config_path],
+        ) == Config(
+            gateways={
+                "testing": GatewayConfig(
+                    connection=MotherDuckConnectionConfig(database="blah"),
+                    state_connection=BigQueryConnectionConfig(),
+                ),
+            },
+            model_defaults=ModelDefaultsConfig(dialect="bigquery"),
+            default_gateway="testing",
+        )
+
+
+def test_load_py_config_env_var_gateway_override(tmp_path_factory):
+    config_path = tmp_path_factory.mktemp("python_config") / "config.py"
+    with open(config_path, "w", encoding="utf-8") as fd:
+        fd.write(
+            """from sqlmesh.core.config import Config, DuckDBConnectionConfig, GatewayConfig, ModelDefaultsConfig
+config = Config(gateways={"duckdb_gateway": GatewayConfig(connection=DuckDBConnectionConfig())}, model_defaults=ModelDefaultsConfig(dialect=''))
+        """
+        )
+    with mock.patch.dict(
+        os.environ,
+        {
+            "SQLMESH__GATEWAYS__DUCKDB_GATEWAY__STATE_CONNECTION__TYPE": "bigquery",
+            "SQLMESH__DEFAULT_GATEWAY": "duckdb_gateway",
+        },
+    ):
+        config = load_config_from_paths(
+            Config,
+            project_paths=[config_path],
+        )
+        assert config == Config(
+            gateways={  # type: ignore
+                "duckdb_gateway": GatewayConfig(
+                    connection=DuckDBConnectionConfig(),
+                    state_connection=BigQueryConnectionConfig(),
+                ),
+            },
+            model_defaults=ModelDefaultsConfig(dialect=""),
+            default_gateway="duckdb_gateway",
+        )
+
+
 def test_load_config_from_python_module_missing_config(tmp_path):
     config_path = tmp_path / "missing_config.py"
     with open(config_path, "w", encoding="utf-8") as fd: