sqlmesh.utils.pydantic

  1from __future__ import annotations
  2
  3import json
  4import typing as t
  5from datetime import tzinfo
  6
  7import pydantic
  8from pydantic import ValidationInfo as ValidationInfo
  9from pydantic.fields import FieldInfo
 10from sqlglot import exp, parse_one
 11from sqlglot.helper import ensure_list
 12from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 13from sqlglot.optimizer.qualify_columns import quote_identifiers
 14
 15from sqlmesh.core import dialect as d
 16from sqlmesh.utils import str_to_bool
 17
 18if t.TYPE_CHECKING:
 19    from sqlglot._typing import E
 20
 21    Model = t.TypeVar("Model", bound="PydanticModel")
 22
 23
 24T = t.TypeVar("T")
 25DEFAULT_ARGS = {"exclude_none": True, "by_alias": True}
 26PRIVATE_FIELDS = "__pydantic_private__"
 27PYDANTIC_MAJOR_VERSION, PYDANTIC_MINOR_VERSION = [int(p) for p in pydantic.__version__.split(".")][
 28    :2
 29]
 30
 31
 32def field_validator(*args: t.Any, **kwargs: t.Any) -> t.Callable[[t.Any], t.Any]:
 33    return pydantic.field_validator(*args, **kwargs)
 34
 35
 36def model_validator(*args: t.Any, **kwargs: t.Any) -> t.Callable[[t.Any], t.Any]:
 37    return pydantic.model_validator(*args, **kwargs)
 38
 39
 40def field_serializer(*args: t.Any, **kwargs: t.Any) -> t.Callable[[t.Any], t.Any]:
 41    return pydantic.field_serializer(*args, **kwargs)
 42
 43
 44def validation_data(info_or_data: t.Any) -> t.Dict[str, t.Any]:
 45    """Safely extract the validated-data dict from a ValidationInfo, dict, or None.
 46
 47    Pydantic 2.13+ sets ValidationInfo.data to None during model_validate_json().
 48    This normalizes all inputs to a dict, returning an empty dict when data is unavailable.
 49    """
 50    if isinstance(info_or_data, dict):
 51        return info_or_data
 52    if info_or_data is not None:
 53        return info_or_data.data or {}
 54    return {}
 55
 56
 57def get_dialect(values: t.Any) -> str:
 58    """Extracts dialect from a dict or pydantic obj, defaulting to the globally set dialect.
 59
 60    Python models allow users to instantiate pydantic models by hand. This is problematic
 61    because the validators kick in with the SQLGLot dialect. To instantiate Pydantic Models used
 62    in python models using the project default dialect, we set a class variable on the model
 63    registry and use that here.
 64    """
 65
 66    from sqlmesh.core.model import model
 67
 68    dialect = validation_data(values).get("dialect")
 69    return model._dialect if dialect is None else dialect  # type: ignore
 70
 71
 72def _expression_encoder(e: exp.Expr) -> str:
 73    return e.meta.get("sql") or e.sql(dialect=e.meta.get("dialect"))
 74
 75
 76AuditQueryTypes = t.Union[exp.Query, d.JinjaQuery]
 77ModelQueryTypes = t.Union[exp.Query, d.JinjaQuery, d.MacroFunc]
 78
 79
 80class PydanticModel(pydantic.BaseModel):
 81    model_config = pydantic.ConfigDict(
 82        # Even though Pydantic v2 kept support for json_encoders, the functionality has been
 83        # crippled badly. Here we need to enumerate all different ways of how sqlglot expressions
 84        # show up in pydantic models.
 85        json_encoders={
 86            exp.Expr: _expression_encoder,
 87            exp.DataType: _expression_encoder,
 88            exp.Tuple: _expression_encoder,
 89            AuditQueryTypes: _expression_encoder,  # type: ignore
 90            ModelQueryTypes: _expression_encoder,  # type: ignore
 91            tzinfo: lambda tz: tz.key,
 92        },
 93        arbitrary_types_allowed=True,
 94        extra="forbid",
 95        protected_namespaces=(),
 96    )
 97
 98    _hash_func_mapping: t.ClassVar[t.Dict[t.Type[t.Any], t.Callable[[t.Any], int]]] = {}
 99
100    def dict(self, **kwargs: t.Any) -> t.Dict[str, t.Any]:
101        kwargs = {**DEFAULT_ARGS, **kwargs}
102        return super().model_dump(**kwargs)  # type: ignore
103
104    def json(
105        self,
106        **kwargs: t.Any,
107    ) -> str:
108        kwargs = {**DEFAULT_ARGS, **kwargs}
109        # Pydantic v2 doesn't support arbitrary arguments for json.dump().
110        if kwargs.pop("sort_keys", False):
111            return json.dumps(super().model_dump(mode="json", **kwargs), sort_keys=True)
112
113        return super().model_dump_json(**kwargs)
114
115    def copy(self: "Model", **kwargs: t.Any) -> "Model":
116        return super().model_copy(**kwargs)
117
118    @property
119    def fields_set(self: "Model") -> t.Set[str]:
120        return self.__pydantic_fields_set__
121
122    @classmethod
123    def parse_obj(cls: t.Type["Model"], obj: t.Any) -> "Model":
124        return super().model_validate(obj)
125
126    @classmethod
127    def parse_raw(cls: t.Type["Model"], b: t.Union[str, bytes], **kwargs: t.Any) -> "Model":
128        return super().model_validate_json(b, **kwargs)
129
130    @classmethod
131    def missing_required_fields(
132        cls: t.Type["PydanticModel"], provided_fields: t.Set[str]
133    ) -> t.Set[str]:
134        return cls.required_fields() - provided_fields
135
136    @classmethod
137    def extra_fields(cls: t.Type["PydanticModel"], provided_fields: t.Set[str]) -> t.Set[str]:
138        return provided_fields - cls.all_fields()
139
140    @classmethod
141    def all_fields(cls: t.Type["PydanticModel"]) -> t.Set[str]:
142        return cls._fields()
143
144    @classmethod
145    def all_field_infos(cls: t.Type["PydanticModel"]) -> t.Dict[str, FieldInfo]:
146        return cls.model_fields
147
148    @classmethod
149    def required_fields(cls: t.Type["PydanticModel"]) -> t.Set[str]:
150        return cls._fields(lambda field: field.is_required())
151
152    @classmethod
153    def _fields(
154        cls: t.Type["PydanticModel"],
155        predicate: t.Callable[[t.Any], bool] = lambda _: True,
156    ) -> t.Set[str]:
157        return {
158            field_info.alias if field_info.alias else field_name
159            for field_name, field_info in cls.all_field_infos().items()
160            if predicate(field_info)
161        }
162
163    def __eq__(self, other: t.Any) -> bool:
164        if (PYDANTIC_MAJOR_VERSION, PYDANTIC_MINOR_VERSION) < (2, 6):
165            if isinstance(other, pydantic.BaseModel):
166                return self.dict() == other.dict()
167            return self.dict() == other
168        return super().__eq__(other)
169
170    def __hash__(self) -> int:
171        if (PYDANTIC_MAJOR_VERSION, PYDANTIC_MINOR_VERSION) < (2, 6):
172            obj = {k: v for k, v in self.__dict__.items() if k in self.all_field_infos()}
173            return hash(self.__class__) + hash(tuple(obj.values()))
174
175        from pydantic._internal._model_construction import make_hash_func  # type: ignore
176
177        if self.__class__ not in PydanticModel._hash_func_mapping:
178            PydanticModel._hash_func_mapping[self.__class__] = make_hash_func(self.__class__)
179
180        return PydanticModel._hash_func_mapping[self.__class__](self)
181
182    def __str__(self) -> str:
183        args = []
184
185        for k, info in self.all_field_infos().items():
186            v = getattr(self, k)
187
188            if type(v) != type(info.default) or v != info.default:
189                args.append(f"{k}: {v}")
190
191        return f"{self.__class__.__name__}<{', '.join(args)}>"
192
193    def __repr__(self) -> str:
194        return str(self)
195
196
197def validate_list_of_strings(v: t.Any) -> t.List[str]:
198    if isinstance(v, exp.Identifier):
199        return [v.name]
200    if isinstance(v, (exp.Tuple, exp.Array)):
201        return [e.name for e in v.expressions]
202    return [i.name if isinstance(i, exp.Identifier) else str(i) for i in v]
203
204
205def validate_string(v: t.Any) -> str:
206    if isinstance(v, exp.Expr):
207        return v.name
208    return str(v)
209
210
211def validate_expression(expression: E, dialect: str) -> E:
212    # this normalizes and quotes identifiers in the given expression according the specified dialect
213    # it also sets expression.meta["dialect"] so that when we serialize for state, the expression is serialized in the correct dialect
214    return _get_field(expression, {"dialect": dialect})  # type: ignore
215
216
217def bool_validator(v: t.Any) -> bool:
218    if isinstance(v, exp.Boolean):
219        return v.this
220    if isinstance(v, exp.Expr):
221        return str_to_bool(v.name)
222    return str_to_bool(str(v or ""))
223
224
225def positive_int_validator(v: t.Any) -> int:
226    if isinstance(v, exp.Expr) and v.is_int:
227        v = int(v.name)
228    if not isinstance(v, int):
229        raise ValueError(f"Invalid num {v}. Value must be an integer value")
230    if v <= 0:
231        raise ValueError(f"Invalid num {v}. Value must be a positive integer")
232    return v
233
234
235def validation_error_message(error: pydantic.ValidationError, base: str) -> str:
236    errors = "\n  ".join(_formatted_validation_errors(error))
237    return f"{base}\n  {errors}"
238
239
240def _formatted_validation_errors(error: pydantic.ValidationError) -> t.List[str]:
241    result = []
242    for e in error.errors():
243        msg = e["msg"]
244        loc: t.Optional[t.Tuple] = e.get("loc")
245        loc_str = ".".join(loc) if loc else None
246        result.append(f"Invalid field '{loc_str}':\n    {msg}" if loc_str else msg)
247    return result
248
249
250def _get_field(
251    v: t.Any,
252    values: t.Any,
253) -> exp.Expr:
254    dialect = get_dialect(values)
255
256    if isinstance(v, exp.Expr):
257        expression = v
258    else:
259        expression = parse_one(v, dialect=dialect)
260
261    expression = exp.column(expression) if isinstance(expression, exp.Identifier) else expression
262    expression = quote_identifiers(
263        normalize_identifiers(expression, dialect=dialect), dialect=dialect
264    )
265    expression.meta["dialect"] = dialect
266
267    return expression
268
269
270def _get_fields(
271    v: t.Any,
272    values: t.Any,
273) -> t.List[exp.Expr]:
274    dialect = get_dialect(values)
275
276    if isinstance(v, (exp.Tuple, exp.Array)):
277        expressions: t.List[exp.Expr] = v.expressions
278    elif isinstance(v, exp.Expr):
279        expressions = [v]
280    else:
281        items: t.List[t.Any] = ensure_list(v)
282        expressions = [
283            parse_one(entry, dialect=dialect) if isinstance(entry, str) else entry  # type: ignore[misc]
284            for entry in items
285        ]
286
287    results = []
288
289    for expr in expressions:
290        results.append(_get_field(expr, values))
291
292    return results
293
294
295def list_of_fields_validator(v: t.Any, values: t.Any) -> t.List[exp.Expr]:
296    return _get_fields(v, values)
297
298
299def column_validator(v: t.Any, values: t.Any) -> exp.Column:
300    expression = _get_field(v, values)
301    if not isinstance(expression, exp.Column):
302        raise ValueError(f"Invalid column {expression}. Value must be a column")
303    return expression
304
305
306def list_of_fields_or_star_validator(
307    v: t.Any, values: t.Any
308) -> t.Union[exp.Star, t.List[exp.Expr]]:
309    expressions = _get_fields(v, values)
310    if len(expressions) == 1 and isinstance(expressions[0], exp.Star):
311        return t.cast(exp.Star, expressions[0])
312    return t.cast(t.List[exp.Expr], expressions)
313
314
315def cron_validator(v: t.Any) -> str:
316    if isinstance(v, exp.Expr):
317        v = v.name
318
319    from croniter import CroniterBadCronError, croniter
320
321    if not isinstance(v, str):
322        raise ValueError(f"Invalid cron expression '{v}'. Value must be a string.")
323
324    try:
325        croniter(v)
326    except CroniterBadCronError:
327        raise ValueError(f"Invalid cron expression '{v}'")
328    return v
329
330
331def get_concrete_types_from_typehint(typehint: type[t.Any]) -> set[type[t.Any]]:
332    concrete_types = set()
333    unpacked = t.get_origin(typehint)
334    if unpacked is None:
335        if type(typehint) == type(type):
336            return {typehint}
337    elif unpacked is t.Union:
338        for item in t.get_args(typehint):
339            if str(item).startswith("typing."):
340                concrete_types |= get_concrete_types_from_typehint(item)
341            else:
342                concrete_types.add(item)
343    else:
344        concrete_types.add(unpacked)
345
346    return concrete_types
347
348
349if t.TYPE_CHECKING:
350    SQLGlotListOfStrings = t.List[str]
351    SQLGlotString = str
352    SQLGlotBool = bool
353    SQLGlotPositiveInt = int
354    SQLGlotColumn = exp.Column
355    SQLGlotListOfFields = t.List[exp.Expr]
356    SQLGlotListOfFieldsOrStar = t.Union[SQLGlotListOfFields, exp.Star]
357    SQLGlotCron = str
358else:
359    from pydantic.functional_validators import BeforeValidator
360
361    SQLGlotListOfStrings = t.Annotated[t.List[str], BeforeValidator(validate_list_of_strings)]
362    SQLGlotString = t.Annotated[str, BeforeValidator(validate_string)]
363    SQLGlotBool = t.Annotated[bool, BeforeValidator(bool_validator)]
364    SQLGlotPositiveInt = t.Annotated[int, BeforeValidator(positive_int_validator)]
365    SQLGlotColumn = t.Annotated[exp.Expr, BeforeValidator(column_validator)]
366    SQLGlotListOfFields = t.Annotated[t.List[exp.Expr], BeforeValidator(list_of_fields_validator)]
367    SQLGlotListOfFieldsOrStar = t.Annotated[
368        t.Union[SQLGlotListOfFields, exp.Star], BeforeValidator(list_of_fields_or_star_validator)
369    ]
370    SQLGlotCron = t.Annotated[str, BeforeValidator(cron_validator)]