sqlmesh.core.config.base

  1from __future__ import annotations
  2
  3import typing as t
  4from enum import Enum, auto
  5
  6from sqlmesh.utils.errors import ConfigError
  7from sqlmesh.utils.pydantic import PydanticModel
  8
  9T = t.TypeVar("T", bound="BaseConfig")
 10
 11
 12class UpdateStrategy(Enum):
 13    """Supported strategies for adding new config to existing config"""
 14
 15    REPLACE = auto()  # Replace with new value
 16    EXTEND = auto()  # Extend list to existing list
 17    KEY_UPDATE = auto()  # Update dict key value with new dict key value
 18    KEY_EXTEND = auto()  # Extend dict key value to existing dict key value
 19    IMMUTABLE = auto()  # Raise if a key tries to change this value
 20    NESTED_UPDATE = auto()  # Recursively updates the nested config
 21
 22
 23def update_field(
 24    old: t.Optional[t.Any],
 25    new: t.Any,
 26    update_strategy: t.Optional[UpdateStrategy] = None,
 27) -> t.Any:
 28    """
 29    Update config field with new config value
 30
 31    Args:
 32        old: The existing config value
 33        new: The new config value
 34        update_strategy: The strategy to use when updating the field
 35
 36
 37        The updated field
 38    """
 39
 40    def _update_pydantic_config(old: BaseConfig, new: BaseConfig) -> PydanticModel:
 41        if type(new) != type(old):
 42            raise ConfigError(
 43                "NESTED_UPDATE behavior requires both values to have the same type. "
 44                f"{type(old)} and {type(new)} were given instead."
 45            )
 46        return old.update_with(new)
 47
 48    if not old:
 49        return new
 50
 51    update_strategy = update_strategy or UpdateStrategy.REPLACE
 52
 53    if update_strategy == UpdateStrategy.IMMUTABLE:
 54        raise ConfigError(f"Cannot modify property: {old}.")
 55
 56    if update_strategy == UpdateStrategy.REPLACE:
 57        return new
 58    if update_strategy == UpdateStrategy.EXTEND:
 59        if not isinstance(old, list) or not isinstance(new, list):
 60            raise ConfigError("EXTEND behavior requires list field.")
 61
 62        return old + new
 63    if update_strategy == UpdateStrategy.KEY_UPDATE:
 64        if not isinstance(old, dict) or not isinstance(new, dict):
 65            raise ConfigError("KEY_UPDATE behavior requires dictionary field.")
 66
 67        combined = old.copy()
 68        combined.update(new)
 69        return combined
 70    if update_strategy == UpdateStrategy.KEY_EXTEND:
 71        if not isinstance(old, dict) or not isinstance(new, dict):
 72            raise ConfigError("KEY_EXTEND behavior requires dictionary field.")
 73
 74        combined = old.copy()
 75        for key, value in new.items():
 76            if not isinstance(value, list):
 77                raise ConfigError("KEY_EXTEND behavior requires list values in dictionary.")
 78
 79            old_value = combined.get(key)
 80            if old_value:
 81                if not isinstance(old_value, list):
 82                    raise ConfigError("KEY_EXTEND behavior requires list values in dictionary.")
 83
 84                combined[key] = old_value + value
 85            else:
 86                combined[key] = value
 87
 88        return combined
 89    if update_strategy == UpdateStrategy.NESTED_UPDATE:
 90        if not isinstance(old, BaseConfig) and not isinstance(old, dict):
 91            raise ConfigError(
 92                f"NESTED_UPDATE behavior requires a config object and a dict of config objects as values. {type(old)} was given instead."
 93            )
 94
 95        if isinstance(old, dict):
 96            for k, pydantic_model in new.items():
 97                if k in old:
 98                    old[k] = _update_pydantic_config(old[k], pydantic_model)
 99                else:
100                    old[k] = pydantic_model
101
102            return old
103        return _update_pydantic_config(old, new)
104
105    raise ConfigError(f"Unknown update strategy {update_strategy}.")
106
107
108class BaseConfig(PydanticModel):
109    """Base configuration functionality for configuration classes."""
110
111    _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {}
112
113    def update_with(self: T, other: t.Union[t.Dict[str, t.Any], T]) -> T:
114        """Updates this instance's fields with the passed in config fields and returns a new instance.
115
116        Args:
117            other: Other configuration.
118
119        Returns:
120            New instance updated with the passed in config fields
121        """
122        if isinstance(other, dict):
123            other = self.__class__(**other)
124
125        updated_fields = {}
126
127        for field in other.fields_set:
128            if field in self.all_field_infos():
129                updated_fields[field] = update_field(
130                    getattr(self, field),
131                    getattr(other, field),
132                    self._FIELD_UPDATE_STRATEGY.get(field),
133                )
134            else:
135                updated_fields[field] = getattr(other, field)
136
137        # Assign each field to trigger assignment validators
138        updated = self.copy()
139        for field, value in updated_fields.items():
140            setattr(updated, field, value)
141
142        return updated