Fix!: Quote identifiers when deserializing model SQL attributes (#2267)

Author: izeigerman

sqlmesh/core/audit/definition.py

@@ -217,7 +217,7 @@ def render_query(
 
         node = t.cast(_Model, node)
         if node.time_column:
-            where = exp.column(node.time_column.column).between(
+            where = node.time_column.column.between(
                 node.convert_to_time_column(start or c.EPOCH, columns_to_types),
                 node.convert_to_time_column(end or c.EPOCH, columns_to_types),
             )

sqlmesh/core/engine_adapter/base.py

@@ -1071,7 +1071,7 @@ def insert_overwrite_by_time_partition(
         time_formatter: t.Callable[
             [TimeLike, t.Optional[t.Dict[str, exp.DataType]]], exp.Expression
         ],
-        time_column: TimeColumn | exp.Column | str,
+        time_column: TimeColumn | exp.Expression | str,
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
         **kwargs: t.Any,
     ) -> None:
@@ -1083,7 +1083,7 @@ def insert_overwrite_by_time_partition(
         if isinstance(time_column, TimeColumn):
             time_column = time_column.column
         where = exp.Between(
-            this=exp.to_column(time_column),
+            this=exp.to_column(time_column) if isinstance(time_column, str) else time_column,
             low=low,
             high=high,
         )
@@ -1939,10 +1939,12 @@ def _to_sql(self, expression: exp.Expression, quote: bool = True, **kwargs: t.An
             **kwargs,
         }
 
+        expression = expression.copy()
+
         if quote:
             quote_identifiers(expression)
 
-        return expression.sql(**sql_gen_kwargs)  # type: ignore
+        return expression.sql(**sql_gen_kwargs, copy=False)  # type: ignore
 
     def _get_data_objects(
         self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None

sqlmesh/core/model/definition.py

@@ -509,12 +509,12 @@ def convert_to_time_column(
             if columns_to_types is None:
                 columns_to_types = self.columns_to_types_or_raise
 
-            if self.time_column.column not in columns_to_types:
+            if self.time_column.column.name not in columns_to_types:
                 raise ConfigError(
-                    f"Time column '{self.time_column.column}' not found in model '{self.name}'."
+                    f"Time column '{self.time_column.column.sql(dialect=self.dialect)}' not found in model '{self.name}'."
                 )
 
-            time_column_type = columns_to_types[self.time_column.column]
+            time_column_type = columns_to_types[self.time_column.column.name]
 
             return to_time_column(time, time_column_type, self.time_column.format)
         return exp.convert(time)
@@ -726,7 +726,7 @@ def _data_hash_values(self) -> t.List[str]:
             data.append(gen(value))
 
         if isinstance(self.kind, IncrementalByTimeRangeKind):
-            data.append(self.kind.time_column.column)
+            data.append(gen(self.kind.time_column.column))
             data.append(self.kind.time_column.format)
         elif isinstance(self.kind, IncrementalByUniqueKeyKind):
             data.extend((gen(k) for k in self.kind.unique_key))

sqlmesh/core/model/kind.py

@@ -6,6 +6,8 @@
 
 from pydantic import Field
 from sqlglot import exp
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+from sqlglot.optimizer.qualify_columns import quote_identifiers
 from sqlglot.time import format_time
 
 from sqlmesh.core import dialect as d
@@ -31,6 +33,10 @@
     from typing_extensions import Annotated, Literal
 
 
+if t.TYPE_CHECKING:
+    MODEL_KIND = t.TypeVar("MODEL_KIND", bound="_ModelKind")
+
+
 class ModelKindMixin:
     @property
     def model_kind_name(self) -> t.Optional[ModelKindName]:
@@ -153,53 +159,76 @@ def data_hash_values(self) -> t.List[t.Optional[str]]:
 
 
 class TimeColumn(PydanticModel):
-    column: str
+    column: exp.Expression
     format: t.Optional[str] = None
 
     @classmethod
     def validator(cls) -> classmethod:
-        def _time_column_validator(v: t.Any) -> TimeColumn:
+        def _time_column_validator(v: t.Any, values: t.Any) -> TimeColumn:
+            values = values if isinstance(values, dict) else values.data
+            dialect = values.get("dialect")
             if isinstance(v, exp.Tuple):
-                kwargs = {
-                    key: v.expressions[i].name
-                    for i, key in enumerate(("column", "format")[: len(v.expressions)])
-                }
-                return TimeColumn(**kwargs)
-
-            if isinstance(v, exp.Expression):
-                return TimeColumn(column=v.name)
+                column_expr = v.expressions[0]
+                column = (
+                    exp.column(column_expr)
+                    if isinstance(column_expr, exp.Identifier)
+                    else column_expr
+                )
+                format = v.expressions[1].name if len(v.expressions) > 1 else None
+            elif isinstance(v, exp.Expression):
+                column = exp.column(v) if isinstance(v, exp.Identifier) else v
+                format = None
+            elif isinstance(v, str):
+                column = d.parse_one(v, dialect=dialect)
+                column.meta.pop("sql")
+                format = None
+            elif isinstance(v, dict):
+                column_raw = v["column"]
+                column = (
+                    d.parse_one(column_raw, dialect=dialect)
+                    if isinstance(column_raw, str)
+                    else column_raw
+                )
+                format = v.get("format")
+            elif isinstance(v, TimeColumn):
+                return v
+            else:
+                raise ConfigError(f"Invalid time_column: '{v}'.")
+
+            column = quote_identifiers(
+                normalize_identifiers(column, dialect=dialect), dialect=dialect
+            )
+            column.meta["dialect"] = dialect
 
-            if isinstance(v, str):
-                return TimeColumn(column=v)
-            return v
+            return TimeColumn(column=column, format=format)
 
         return field_validator("time_column", mode="before")(_time_column_validator)
 
     @field_validator("column", mode="before")
     @classmethod
-    def _column_validator(cls, v: str) -> str:
+    def _column_validator(cls, v: t.Union[str, exp.Expression]) -> exp.Expression:
         if not v:
             raise ConfigError("Time Column cannot be empty.")
+        if isinstance(v, str):
+            return exp.to_column(v)
         return v
 
     @property
-    def expression(self) -> exp.Column | exp.Tuple:
+    def expression(self) -> exp.Expression:
         """Convert this pydantic model into a time_column SQLGlot expression."""
-        column = exp.to_column(self.column)
         if not self.format:
-            return column
+            return self.column
 
-        return exp.Tuple(expressions=[column, exp.Literal.string(self.format)])
+        return exp.Tuple(expressions=[self.column, exp.Literal.string(self.format)])
 
-    def to_expression(self, dialect: str) -> exp.Column | exp.Tuple:
+    def to_expression(self, dialect: str) -> exp.Expression:
         """Convert this pydantic model into a time_column SQLGlot expression."""
-        column = exp.to_column(self.column)
         if not self.format:
-            return column
+            return self.column
 
         return exp.Tuple(
             expressions=[
-                column,
+                self.column,
                 exp.Literal.string(
                     format_time(self.format, d.Dialect.get_or_raise(dialect).INVERSE_TIME_MAPPING)
                 ),
@@ -211,6 +240,7 @@ def to_property(self, dialect: str = "") -> exp.Property:
 
 
 class _Incremental(_ModelKind):
+    dialect: str = ""
     batch_size: t.Optional[SQLGlotPositiveInt] = None
     lookback: t.Optional[SQLGlotPositiveInt] = None
     forward_only: SQLGlotBool = False
@@ -335,6 +365,7 @@ class FullKind(_ModelKind):
 
 
 class _SCDType2Kind(_ModelKind):
+    dialect: str = ""
     unique_key: SQLGlotListOfFields
     valid_from_name: SQLGlotString = "valid_from"
     valid_to_name: SQLGlotString = "valid_to"
@@ -344,6 +375,16 @@ class _SCDType2Kind(_ModelKind):
     forward_only: SQLGlotBool = True
     disable_restatement: SQLGlotBool = True
 
+    @field_validator("time_data_type", mode="before")
+    @classmethod
+    def _time_data_type_validator(
+        cls, v: t.Union[str, exp.Expression], values: t.Any
+    ) -> exp.Expression:
+        values = values if isinstance(values, dict) else values.data
+        if isinstance(v, exp.Expression) and not isinstance(v, exp.DataType):
+            v = v.name
+        return exp.DataType.build(v, dialect=values.get("dialect"))
+
     @property
     def managed_columns(self) -> t.Dict[str, exp.DataType]:
         return {
@@ -425,20 +466,16 @@ def _model_kind_validator(cls: t.Type, v: t.Any, values: t.Dict[str, t.Any]) ->
             if isinstance(v, d.ModelKind)
             else v
         )
-        time_data_type = props.pop("time_data_type", None)
-        if isinstance(time_data_type, exp.Expression) and not isinstance(
-            time_data_type, exp.DataType
-        ):
-            time_data_type = time_data_type.name
-        if time_data_type:
-            props["time_data_type"] = exp.DataType.build(time_data_type, dialect=dialect)
         name = v.this if isinstance(v, d.ModelKind) else props.get("name")
         # We want to ensure whatever name is provided to construct the class is the same name that will be
         # found inside the class itself in order to avoid a change during plan/apply for legacy aliases.
         # Ex: Pass in `SCD_TYPE_2` then we want to ensure we get `SCD_TYPE_2` as the kind name
         # instead of `SCD_TYPE_2_BY_TIME`.
         props["name"] = name
-        return model_kind_type_from_name(name)(**props)
+        kind_type = model_kind_type_from_name(name)
+        if "dialect" in kind_type.all_fields() and props.get("dialect") is None:
+            props["dialect"] = dialect
+        return kind_type(**props)
 
     name = (v.name if isinstance(v, exp.Expression) else str(v)).upper()
     return model_kind_type_from_name(name)(name=name)  # type: ignore

sqlmesh/core/model/meta.py

@@ -267,19 +267,6 @@ def _kind_validator(cls, values: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
             for field in ("partitioned_by_", "clustered_by"):
                 if values.get(field) and not kind.is_materialized:
                     raise ValueError(f"{field} field cannot be set for {kind} models")
-
-            dialect = values.get("dialect")
-
-            if hasattr(kind, "time_column"):
-                kind.time_column.column = normalize_identifiers(
-                    kind.time_column.column, dialect=dialect
-                ).name
-
-            if hasattr(kind, "unique_key"):
-                kind.unique_key = [
-                    normalize_identifiers(key, dialect=dialect) for key in kind.unique_key
-                ]
-
         return values
 
     @property
@@ -298,9 +285,9 @@ def unique_key(self) -> t.List[exp.Expression]:
     @property
     def partitioned_by(self) -> t.List[exp.Expression]:
         if self.time_column and self.time_column.column not in [
-            col.name for col in self._partition_by_columns
+            col for col in self._partition_by_columns
         ]:
-            return [*[exp.to_column(self.time_column.column)], *self.partitioned_by_]
+            return [self.time_column.column, *self.partitioned_by_]
         return self.partitioned_by_
 
     @property

sqlmesh/dbt/model.py

@@ -184,7 +184,7 @@ def model_kind(self, context: DbtContext) -> ModelKind:
         if materialization == Materialization.VIEW:
             return ViewKind()
         if materialization == Materialization.INCREMENTAL:
-            incremental_kwargs = {}
+            incremental_kwargs = {"dialect": context.dialect}
             for field in ("batch_size", "lookback", "forward_only", "disable_restatement"):
                 field_val = getattr(self, field, None) or self.meta.get(field, None)
                 if field_val:
@@ -243,6 +243,7 @@ def model_kind(self, context: DbtContext) -> ModelKind:
                     f"{self.canonical_name(context)}: SQLMesh snapshot strategy is required for snapshot materialization."
                 )
             shared_kwargs = {
+                "dialect": context.dialect,
                 "unique_key": self.unique_key,
                 "invalidate_hard_deletes": self.invalidate_hard_deletes,
                 "valid_from_name": "dbt_valid_from",

sqlmesh/migrations/v0044_quote_identifiers_in_model_attributes.py

@@ -0,0 +1,5 @@
+"""Quoted identifiers in model SQL attributes."""
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    pass

sqlmesh/utils/pydantic.py

@@ -10,6 +10,7 @@
 from sqlglot import exp, parse_one
 from sqlglot.helper import ensure_list
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+from sqlglot.optimizer.qualify_columns import quote_identifiers
 
 from sqlmesh.core import dialect as d
 from sqlmesh.utils import str_to_bool
@@ -309,10 +310,8 @@ def _get_fields(
     results = []
 
     for expr in expressions:
-        expr = normalize_identifiers(
-            exp.column(expr) if isinstance(expr, exp.Identifier) else expr,
-            dialect=dialect,
-        )
+        expr = exp.column(expr) if isinstance(expr, exp.Identifier) else expr
+        expr = quote_identifiers(normalize_identifiers(expr, dialect=dialect), dialect=dialect)
         expr.meta["dialect"] = dialect
         results.append(expr)