Introduce config env_vars (#1047)

Author: vchan

sqlmesh/core/config/__init__.py

@@ -10,7 +10,7 @@
     SparkConnectionConfig,
 )
 from sqlmesh.core.config.gateway import GatewayConfig
-from sqlmesh.core.config.loader import load_config_from_paths
+from sqlmesh.core.config.loader import load_config_from_paths, load_config_from_yaml
 from sqlmesh.core.config.model import ModelDefaultsConfig
 from sqlmesh.core.config.root import Config
 from sqlmesh.core.config.scheduler import AirflowSchedulerConfig, BuiltInSchedulerConfig

sqlmesh/core/config/loader.py

@@ -13,7 +13,10 @@
 
 
 def load_config_from_paths(
-    *paths: Path, config_name: str = "config", load_from_env: bool = True
+    *paths: Path,
+    config_name: str = "config",
+    load_from_env: bool = True,
+    config_defaults: Config | None = None,
 ) -> Config:
     config: t.Optional[Config] = None
 
@@ -34,7 +37,7 @@ def load_config_from_paths(
         if extension in ("yml", "yaml"):
             if config_name != "config":
                 raise ConfigError(
-                    f"YAML configs do not support multiple configs. Use Python instead."
+                    "YAML configs do not support multiple configs. Use Python instead."
                 )
             next_config = load_config_from_yaml(path)
         elif extension == "py":
@@ -54,6 +57,9 @@ def load_config_from_paths(
     if load_from_env:
         config = config.update_with(load_config_from_env())
 
+    if config_defaults:
+        config = config_defaults.update_with(config)
+
     return config
 
 

sqlmesh/core/config/root.py

@@ -55,6 +55,7 @@ class Config(BaseConfig):
     users: t.List[User] = []
     model_defaults: ModelDefaultsConfig = ModelDefaultsConfig()
     loader: t.Type[Loader] = SqlMeshLoader
+    env_vars: t.Dict[str, str] = {}
 
     _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
         "gateways": UpdateStrategy.KEY_UPDATE,

sqlmesh/core/context.py

@@ -47,7 +47,7 @@
 from sqlmesh.core import constants as c
 from sqlmesh.core._typing import NotificationTarget
 from sqlmesh.core.audit import Audit
-from sqlmesh.core.config import Config, load_config_from_paths
+from sqlmesh.core.config import Config, load_config_from_paths, load_config_from_yaml
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.context_diff import ContextDiff
 from sqlmesh.core.dialect import (
@@ -75,7 +75,7 @@
 from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.core.test import get_all_model_tests, run_model_tests, run_tests
 from sqlmesh.core.user import User
-from sqlmesh.utils import UniqueKeyDict, sys_path
+from sqlmesh.utils import UniqueKeyDict, env_vars, sys_path
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import TimeLike, yesterday_ds
 from sqlmesh.utils.errors import (
@@ -1082,15 +1082,23 @@ def _load_configs(
         if isinstance(config, Config):
             return {path: config for path in paths}
 
-        lookup_paths = [self.sqlmesh_path / "config.yml", self.sqlmesh_path / "config.yaml"]
-
-        return {
-            path: load_config_from_paths(
-                *(lookup_paths + [path / "config.py", path / "config.yml", path / "config.yaml"]),
-                config_name=config,
-            )
-            for path in paths
-        }
+        config_defaults = None
+        for path in (self.sqlmesh_path / "config.yml", self.sqlmesh_path / "config.yaml"):
+            if path.exists():
+                config_defaults = load_config_from_yaml(path)
+                break
+
+        with env_vars(config_defaults.env_vars if config_defaults else {}):
+            return {
+                path: load_config_from_paths(
+                    path / "config.py",
+                    path / "config.yml",
+                    path / "config.yaml",
+                    config_name=config,
+                    config_defaults=config_defaults,
+                )
+                for path in paths
+            }
 
     def _run_janitor(self) -> None:
         expired_environments = self.state_sync.delete_expired_environments()

sqlmesh/utils/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import importlib
+import os
 import re
 import sys
 import time
@@ -218,3 +219,16 @@ class classproperty(property):
 
     def __get__(self, obj: t.Any, owner: t.Any = None) -> t.Any:
         return classmethod(self.fget).__get__(None, owner)()  # type: ignore
+
+
+@contextmanager
+def env_vars(environ: dict[str, str]) -> t.Iterator[None]:
+    """A context manager to temporarily modify environment variables."""
+    old_environ = os.environ.copy()
+    os.environ.update(environ)
+
+    try:
+        yield
+    finally:
+        os.environ.clear()
+        os.environ.update(old_environ)