sqlmesh.core.model.kind

  1from __future__ import annotations
  2
  3import sys
  4import typing as t
  5from enum import Enum
  6
  7from pydantic import Field
  8from sqlglot import exp
  9from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 10from sqlglot.optimizer.qualify_columns import quote_identifiers
 11from sqlglot.optimizer.simplify import gen
 12from sqlglot.time import format_time
 13
 14from sqlmesh.core import dialect as d
 15from sqlmesh.core.model.common import parse_properties
 16from sqlmesh.core.model.seed import CsvSettings
 17from sqlmesh.utils.errors import ConfigError
 18from sqlmesh.utils.pydantic import (
 19    PydanticModel,
 20    SQLGlotBool,
 21    SQLGlotColumn,
 22    SQLGlotListOfColumnsOrStar,
 23    SQLGlotListOfFields,
 24    SQLGlotPositiveInt,
 25    SQLGlotString,
 26    column_validator,
 27    field_validator,
 28    field_validator_v1_args,
 29    get_dialect,
 30)
 31
 32if sys.version_info >= (3, 9):
 33    from typing import Annotated, Literal
 34else:
 35    from typing_extensions import Annotated, Literal
 36
 37
 38if t.TYPE_CHECKING:
 39    MODEL_KIND = t.TypeVar("MODEL_KIND", bound="_ModelKind")
 40
 41
 42class ModelKindMixin:
 43    @property
 44    def model_kind_name(self) -> t.Optional[ModelKindName]:
 45        """Returns the model kind name."""
 46        raise NotImplementedError
 47
 48    @property
 49    def is_incremental_by_time_range(self) -> bool:
 50        return self.model_kind_name == ModelKindName.INCREMENTAL_BY_TIME_RANGE
 51
 52    @property
 53    def is_incremental_by_unique_key(self) -> bool:
 54        return self.model_kind_name == ModelKindName.INCREMENTAL_BY_UNIQUE_KEY
 55
 56    @property
 57    def is_incremental_unmanaged(self) -> bool:
 58        return self.model_kind_name == ModelKindName.INCREMENTAL_UNMANAGED
 59
 60    @property
 61    def is_incremental(self) -> bool:
 62        return (
 63            self.is_incremental_by_time_range
 64            or self.is_incremental_by_unique_key
 65            or self.is_incremental_unmanaged
 66            or self.is_scd_type_2
 67        )
 68
 69    @property
 70    def is_full(self) -> bool:
 71        return self.model_kind_name == ModelKindName.FULL
 72
 73    @property
 74    def is_view(self) -> bool:
 75        return self.model_kind_name == ModelKindName.VIEW
 76
 77    @property
 78    def is_embedded(self) -> bool:
 79        return self.model_kind_name == ModelKindName.EMBEDDED
 80
 81    @property
 82    def is_seed(self) -> bool:
 83        return self.model_kind_name == ModelKindName.SEED
 84
 85    @property
 86    def is_external(self) -> bool:
 87        return self.model_kind_name == ModelKindName.EXTERNAL
 88
 89    @property
 90    def is_scd_type_2(self) -> bool:
 91        return self.model_kind_name in {
 92            ModelKindName.SCD_TYPE_2,
 93            ModelKindName.SCD_TYPE_2_BY_TIME,
 94            ModelKindName.SCD_TYPE_2_BY_COLUMN,
 95        }
 96
 97    @property
 98    def is_scd_type_2_by_time(self) -> bool:
 99        return self.model_kind_name in {ModelKindName.SCD_TYPE_2, ModelKindName.SCD_TYPE_2_BY_TIME}
100
101    @property
102    def is_scd_type_2_by_column(self) -> bool:
103        return self.model_kind_name == ModelKindName.SCD_TYPE_2_BY_COLUMN
104
105    @property
106    def is_symbolic(self) -> bool:
107        """A symbolic model is one that doesn't execute at all."""
108        return self.model_kind_name in (ModelKindName.EMBEDDED, ModelKindName.EXTERNAL)
109
110    @property
111    def is_materialized(self) -> bool:
112        return not (self.is_symbolic or self.is_view)
113
114    @property
115    def only_execution_time(self) -> bool:
116        """Whether or not this model only cares about execution time to render."""
117        return self.is_view or self.is_full or self.is_scd_type_2
118
119    @property
120    def full_history_restatement_only(self) -> bool:
121        """Whether or not this model only supports restatement of full history."""
122        return (
123            self.is_incremental_unmanaged or self.is_incremental_by_unique_key or self.is_scd_type_2
124        )
125
126
127class ModelKindName(str, ModelKindMixin, Enum):
128    """The kind of model, determining how this data is computed and stored in the warehouse."""
129
130    INCREMENTAL_BY_TIME_RANGE = "INCREMENTAL_BY_TIME_RANGE"
131    INCREMENTAL_BY_UNIQUE_KEY = "INCREMENTAL_BY_UNIQUE_KEY"
132    INCREMENTAL_UNMANAGED = "INCREMENTAL_UNMANAGED"
133    FULL = "FULL"
134    # Legacy alias to SCD Type 2 By Time
135    # Only used for Parsing and mapping name to SCD Type 2 By Time
136    SCD_TYPE_2 = "SCD_TYPE_2"
137    SCD_TYPE_2_BY_TIME = "SCD_TYPE_2_BY_TIME"
138    SCD_TYPE_2_BY_COLUMN = "SCD_TYPE_2_BY_COLUMN"
139    VIEW = "VIEW"
140    EMBEDDED = "EMBEDDED"
141    SEED = "SEED"
142    EXTERNAL = "EXTERNAL"
143
144    @property
145    def model_kind_name(self) -> t.Optional[ModelKindName]:
146        return self
147
148    def __str__(self) -> str:
149        return self.name
150
151    def __repr__(self) -> str:
152        return str(self)
153
154
155class _ModelKind(PydanticModel, ModelKindMixin):
156    name: ModelKindName
157
158    @property
159    def model_kind_name(self) -> t.Optional[ModelKindName]:
160        return self.name
161
162    def to_expression(self, **kwargs: t.Any) -> d.ModelKind:
163        return d.ModelKind(this=self.name.value.upper(), **kwargs)
164
165    @property
166    def data_hash_values(self) -> t.List[t.Optional[str]]:
167        return [self.name.value]
168
169    @property
170    def metadata_hash_values(self) -> t.List[t.Optional[str]]:
171        return []
172
173
174class TimeColumn(PydanticModel):
175    column: exp.Expression
176    format: t.Optional[str] = None
177
178    @classmethod
179    def validator(cls) -> classmethod:
180        def _time_column_validator(v: t.Any, values: t.Any) -> TimeColumn:
181            dialect = get_dialect(values)
182
183            if isinstance(v, exp.Tuple):
184                column_expr = v.expressions[0]
185                column = (
186                    exp.column(column_expr)
187                    if isinstance(column_expr, exp.Identifier)
188                    else column_expr
189                )
190                format = v.expressions[1].name if len(v.expressions) > 1 else None
191            elif isinstance(v, exp.Expression):
192                column = exp.column(v) if isinstance(v, exp.Identifier) else v
193                format = None
194            elif isinstance(v, str):
195                column = d.parse_one(v, dialect=dialect)
196                column.meta.pop("sql")
197                format = None
198            elif isinstance(v, dict):
199                column_raw = v["column"]
200                column = (
201                    d.parse_one(column_raw, dialect=dialect)
202                    if isinstance(column_raw, str)
203                    else column_raw
204                )
205                format = v.get("format")
206            elif isinstance(v, TimeColumn):
207                column = v.column
208                format = v.format
209            else:
210                raise ConfigError(f"Invalid time_column: '{v}'.")
211
212            column = quote_identifiers(
213                normalize_identifiers(column, dialect=dialect), dialect=dialect
214            )
215            column.meta["dialect"] = dialect
216
217            return TimeColumn(column=column, format=format)
218
219        return field_validator("time_column", mode="before")(_time_column_validator)
220
221    @field_validator("column", mode="before")
222    @classmethod
223    def _column_validator(cls, v: t.Union[str, exp.Expression]) -> exp.Expression:
224        if not v:
225            raise ConfigError("Time Column cannot be empty.")
226        if isinstance(v, str):
227            return exp.to_column(v)
228        return v
229
230    @property
231    def expression(self) -> exp.Expression:
232        """Convert this pydantic model into a time_column SQLGlot expression."""
233        if not self.format:
234            return self.column
235
236        return exp.Tuple(expressions=[self.column, exp.Literal.string(self.format)])
237
238    def to_expression(self, dialect: str) -> exp.Expression:
239        """Convert this pydantic model into a time_column SQLGlot expression."""
240        if not self.format:
241            return self.column
242
243        return exp.Tuple(
244            expressions=[
245                self.column,
246                exp.Literal.string(
247                    format_time(self.format, d.Dialect.get_or_raise(dialect).INVERSE_TIME_MAPPING)
248                ),
249            ]
250        )
251
252    def to_property(self, dialect: str = "") -> exp.Property:
253        return exp.Property(this="time_column", value=self.to_expression(dialect))
254
255
256def _kind_dialect_validator(cls: t.Type, v: t.Optional[str]) -> str:
257    if v is None:
258        return get_dialect({})
259    return v
260
261
262kind_dialect_validator = field_validator("dialect", mode="before", always=True)(
263    _kind_dialect_validator
264)
265
266
267class _Incremental(_ModelKind):
268    dialect: t.Optional[str] = Field(None, validate_default=True)
269    batch_size: t.Optional[SQLGlotPositiveInt] = None
270    batch_concurrency: t.Optional[SQLGlotPositiveInt] = None
271    lookback: t.Optional[SQLGlotPositiveInt] = None
272    forward_only: SQLGlotBool = False
273    disable_restatement: SQLGlotBool = False
274
275    _dialect_validator = kind_dialect_validator
276
277    @property
278    def data_hash_values(self) -> t.List[t.Optional[str]]:
279        return [
280            *super().data_hash_values,
281            self.dialect,
282            str(self.lookback) if self.lookback is not None else None,
283        ]
284
285    @property
286    def metadata_hash_values(self) -> t.List[t.Optional[str]]:
287        return [
288            *super().metadata_hash_values,
289            str(self.batch_size) if self.batch_size is not None else None,
290            str(self.forward_only),
291            str(self.disable_restatement),
292        ]
293
294
295class IncrementalByTimeRangeKind(_Incremental):
296    name: Literal[ModelKindName.INCREMENTAL_BY_TIME_RANGE] = ModelKindName.INCREMENTAL_BY_TIME_RANGE
297    time_column: TimeColumn
298
299    _time_column_validator = TimeColumn.validator()
300
301    def to_expression(self, dialect: str = "", **kwargs: t.Any) -> d.ModelKind:
302        return super().to_expression(expressions=[self.time_column.to_property(dialect)])
303
304    @property
305    def data_hash_values(self) -> t.List[t.Optional[str]]:
306        return [*super().data_hash_values, gen(self.time_column.column), self.time_column.format]
307
308
309class IncrementalByUniqueKeyKind(_Incremental):
310    name: Literal[ModelKindName.INCREMENTAL_BY_UNIQUE_KEY] = ModelKindName.INCREMENTAL_BY_UNIQUE_KEY
311    unique_key: SQLGlotListOfFields
312    when_matched: t.Optional[exp.When] = None
313    batch_concurrency: Literal[1] = 1
314
315    @field_validator("when_matched", mode="before")
316    @field_validator_v1_args
317    def _when_matched_validator(
318        cls, v: t.Optional[t.Union[exp.When, str]], values: t.Dict[str, t.Any]
319    ) -> t.Optional[exp.When]:
320        def replace_table_references(expression: exp.Expression) -> exp.Expression:
321            from sqlmesh.core.engine_adapter.base import (
322                MERGE_SOURCE_ALIAS,
323                MERGE_TARGET_ALIAS,
324            )
325
326            if isinstance(expression, exp.Column):
327                if expression.table.lower() == "target":
328                    expression.set(
329                        "table",
330                        exp.to_identifier(MERGE_TARGET_ALIAS),
331                    )
332                elif expression.table.lower() == "source":
333                    expression.set(
334                        "table",
335                        exp.to_identifier(MERGE_SOURCE_ALIAS),
336                    )
337            return expression
338
339        if isinstance(v, str):
340            return t.cast(exp.When, d.parse_one(v, into=exp.When, dialect=get_dialect(values)))
341
342        if not v:
343            return v
344
345        return t.cast(exp.When, v.transform(replace_table_references))
346
347    @property
348    def data_hash_values(self) -> t.List[t.Optional[str]]:
349        return [
350            *super().data_hash_values,
351            *(gen(k) for k in self.unique_key),
352            gen(self.when_matched) if self.when_matched is not None else None,
353        ]
354
355
356class IncrementalUnmanagedKind(_ModelKind):
357    name: Literal[ModelKindName.INCREMENTAL_UNMANAGED] = ModelKindName.INCREMENTAL_UNMANAGED
358    insert_overwrite: SQLGlotBool = False
359    forward_only: SQLGlotBool = True
360    disable_restatement: SQLGlotBool = True
361
362    @property
363    def data_hash_values(self) -> t.List[t.Optional[str]]:
364        return [*super().data_hash_values, str(self.insert_overwrite)]
365
366    @property
367    def metadata_hash_values(self) -> t.List[t.Optional[str]]:
368        return [*super().metadata_hash_values, str(self.forward_only)]
369
370
371class ViewKind(_ModelKind):
372    name: Literal[ModelKindName.VIEW] = ModelKindName.VIEW
373    materialized: SQLGlotBool = False
374
375    @property
376    def data_hash_values(self) -> t.List[t.Optional[str]]:
377        return [*super().data_hash_values, str(self.materialized)]
378
379
380class SeedKind(_ModelKind):
381    name: Literal[ModelKindName.SEED] = ModelKindName.SEED
382    path: SQLGlotString
383    batch_size: SQLGlotPositiveInt = 1000
384    csv_settings: t.Optional[CsvSettings] = None
385
386    @field_validator("csv_settings", mode="before")
387    @classmethod
388    def _parse_csv_settings(cls, v: t.Any) -> t.Optional[CsvSettings]:
389        if v is None or isinstance(v, CsvSettings):
390            return v
391        if isinstance(v, exp.Expression):
392            tuple_exp = parse_properties(cls, v, {})
393            if not tuple_exp:
394                return None
395            return CsvSettings(**{e.left.name: e.right for e in tuple_exp.expressions})
396        if isinstance(v, dict):
397            return CsvSettings(**v)
398        return v
399
400    def to_expression(self, **kwargs: t.Any) -> d.ModelKind:
401        """Convert the seed kind into a SQLGlot expression."""
402        return super().to_expression(
403            expressions=[
404                exp.Property(this=exp.Var(this="path"), value=exp.Literal.string(self.path)),
405                exp.Property(
406                    this=exp.Var(this="batch_size"),
407                    value=exp.Literal.number(self.batch_size),
408                ),
409            ],
410        )
411
412    @property
413    def data_hash_values(self) -> t.List[t.Optional[str]]:
414        return [
415            *super().data_hash_values,
416            *(self.csv_settings or CsvSettings()).dict().values(),
417        ]
418
419    @property
420    def metadata_hash_values(self) -> t.List[t.Optional[str]]:
421        return [*super().metadata_hash_values, str(self.batch_size)]
422
423
424class FullKind(_ModelKind):
425    name: Literal[ModelKindName.FULL] = ModelKindName.FULL
426
427
428class _SCDType2Kind(_ModelKind):
429    dialect: t.Optional[str] = Field(None, validate_default=True)
430    unique_key: SQLGlotListOfFields
431    valid_from_name: SQLGlotColumn = Field(exp.column("valid_from"), validate_default=True)
432    valid_to_name: SQLGlotColumn = Field(exp.column("valid_to"), validate_default=True)
433    invalidate_hard_deletes: SQLGlotBool = False
434    time_data_type: exp.DataType = Field(exp.DataType.build("TIMESTAMP"), validate_default=True)
435
436    forward_only: SQLGlotBool = True
437    disable_restatement: SQLGlotBool = True
438
439    _dialect_validator = kind_dialect_validator
440    # Remove once Pydantic 1 is deprecated
441    _always_validate_column = field_validator(
442        "valid_from_name", "valid_to_name", mode="before", always=True
443    )(column_validator)
444
445    # always=True can be removed once Pydantic 1 is deprecated
446    @field_validator("time_data_type", mode="before", always=True)
447    @classmethod
448    def _time_data_type_validator(
449        cls, v: t.Union[str, exp.Expression], values: t.Any
450    ) -> exp.Expression:
451        if isinstance(v, exp.Expression) and not isinstance(v, exp.DataType):
452            v = v.name
453        dialect = get_dialect(values)
454        data_type = exp.DataType.build(v, dialect=dialect)
455        data_type.meta["dialect"] = dialect
456        return data_type
457
458    @property
459    def managed_columns(self) -> t.Dict[str, exp.DataType]:
460        return {
461            self.valid_from_name.name: self.time_data_type,
462            self.valid_to_name.name: self.time_data_type,
463        }
464
465    @property
466    def data_hash_values(self) -> t.List[t.Optional[str]]:
467        return [
468            *super().data_hash_values,
469            self.dialect,
470            *(gen(k) for k in self.unique_key),
471            gen(self.valid_from_name),
472            gen(self.valid_to_name),
473            str(self.invalidate_hard_deletes),
474            gen(self.time_data_type),
475        ]
476
477    @property
478    def metadata_hash_values(self) -> t.List[t.Optional[str]]:
479        return [
480            *super().metadata_hash_values,
481            str(self.forward_only),
482            str(self.disable_restatement),
483        ]
484
485
486class SCDType2ByTimeKind(_SCDType2Kind):
487    name: Literal[ModelKindName.SCD_TYPE_2, ModelKindName.SCD_TYPE_2_BY_TIME] = (
488        ModelKindName.SCD_TYPE_2_BY_TIME
489    )
490    updated_at_name: SQLGlotColumn = Field(exp.column("updated_at"), validate_default=True)
491    updated_at_as_valid_from: SQLGlotBool = False
492
493    # Remove once Pydantic 1 is deprecated
494    _always_validate_updated_at = field_validator("updated_at_name", mode="before", always=True)(
495        column_validator
496    )
497
498    @property
499    def data_hash_values(self) -> t.List[t.Optional[str]]:
500        return [
501            *super().data_hash_values,
502            gen(self.updated_at_name),
503            str(self.updated_at_as_valid_from),
504        ]
505
506
507class SCDType2ByColumnKind(_SCDType2Kind):
508    name: Literal[ModelKindName.SCD_TYPE_2_BY_COLUMN] = ModelKindName.SCD_TYPE_2_BY_COLUMN
509    columns: SQLGlotListOfColumnsOrStar
510    execution_time_as_valid_from: SQLGlotBool = False
511
512    @property
513    def data_hash_values(self) -> t.List[t.Optional[str]]:
514        columns_sql = (
515            [gen(c) for c in self.columns]
516            if isinstance(self.columns, list)
517            else [gen(self.columns)]
518        )
519        return [*super().data_hash_values, *columns_sql, str(self.execution_time_as_valid_from)]
520
521
522class EmbeddedKind(_ModelKind):
523    name: Literal[ModelKindName.EMBEDDED] = ModelKindName.EMBEDDED
524
525
526class ExternalKind(_ModelKind):
527    name: Literal[ModelKindName.EXTERNAL] = ModelKindName.EXTERNAL
528
529
530ModelKind = Annotated[
531    t.Union[
532        EmbeddedKind,
533        ExternalKind,
534        FullKind,
535        IncrementalByTimeRangeKind,
536        IncrementalByUniqueKeyKind,
537        IncrementalUnmanagedKind,
538        SeedKind,
539        ViewKind,
540        SCDType2ByTimeKind,
541        SCDType2ByColumnKind,
542    ],
543    Field(discriminator="name"),
544]
545
546MODEL_KIND_NAME_TO_TYPE: t.Dict[str, t.Type[ModelKind]] = {
547    ModelKindName.EMBEDDED: EmbeddedKind,
548    ModelKindName.EXTERNAL: ExternalKind,
549    ModelKindName.FULL: FullKind,
550    ModelKindName.INCREMENTAL_BY_TIME_RANGE: IncrementalByTimeRangeKind,
551    ModelKindName.INCREMENTAL_BY_UNIQUE_KEY: IncrementalByUniqueKeyKind,
552    ModelKindName.INCREMENTAL_UNMANAGED: IncrementalUnmanagedKind,
553    ModelKindName.SEED: SeedKind,
554    ModelKindName.VIEW: ViewKind,
555    ModelKindName.SCD_TYPE_2: SCDType2ByTimeKind,
556    ModelKindName.SCD_TYPE_2_BY_TIME: SCDType2ByTimeKind,
557    ModelKindName.SCD_TYPE_2_BY_COLUMN: SCDType2ByColumnKind,
558}
559
560
561def model_kind_type_from_name(name: t.Optional[str]) -> t.Type[ModelKind]:
562    klass = MODEL_KIND_NAME_TO_TYPE.get(name) if name else None
563    if not klass:
564        raise ConfigError(f"Invalid model kind '{name}'")
565    return t.cast(t.Type[ModelKind], klass)
566
567
568@field_validator_v1_args
569def _model_kind_validator(cls: t.Type, v: t.Any, values: t.Dict[str, t.Any]) -> ModelKind:
570    dialect = get_dialect(values)
571
572    if isinstance(v, _ModelKind):
573        return t.cast(ModelKind, v)
574
575    if isinstance(v, (d.ModelKind, dict)):
576        props = (
577            {prop.name: prop.args.get("value") for prop in v.expressions}
578            if isinstance(v, d.ModelKind)
579            else v
580        )
581        name = v.this if isinstance(v, d.ModelKind) else props.get("name")
582        # We want to ensure whatever name is provided to construct the class is the same name that will be
583        # found inside the class itself in order to avoid a change during plan/apply for legacy aliases.
584        # Ex: Pass in `SCD_TYPE_2` then we want to ensure we get `SCD_TYPE_2` as the kind name
585        # instead of `SCD_TYPE_2_BY_TIME`.
586        props["name"] = name
587        kind_type = model_kind_type_from_name(name)
588        if "dialect" in kind_type.all_fields() and props.get("dialect") is None:
589            props["dialect"] = dialect
590        return kind_type(**props)
591
592    name = (v.name if isinstance(v, exp.Expression) else str(v)).upper()
593    return model_kind_type_from_name(name)(name=name)  # type: ignore
594
595
596model_kind_validator = field_validator("kind", mode="before")(_model_kind_validator)