Feat: accept expressions in unique key (#1528)

closes #1516

Author: tobymao

sqlmesh/core/dialect.py

@@ -814,6 +814,19 @@ def find_tables(expression: exp.Expression, dialect: DialectType = None) -> t.Se
     }
 
 
+def add_table(node: exp.Expression, table: str) -> exp.Expression:
+    """Add a table to all columns in an expression."""
+
+    def _transform(node: exp.Expression) -> exp.Expression:
+        if isinstance(node, exp.Column) and not node.table:
+            return exp.column(node.this, table=table)
+        if isinstance(node, exp.Identifier):
+            return exp.column(node, table=table)
+        return node
+
+    return node.transform(_transform)
+
+
 def transform_values(
     values: t.Tuple[t.Any, ...], columns_to_types: t.Dict[str, exp.DataType]
 ) -> t.Iterator[t.Any]:

sqlmesh/core/engine_adapter/base.py

@@ -25,7 +25,7 @@
 from sqlglot.helper import ensure_list
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 
-from sqlmesh.core.dialect import select_from_values_for_batch_range
+from sqlmesh.core.dialect import add_table, select_from_values_for_batch_range
 from sqlmesh.core.engine_adapter.shared import DataObject, TransactionType
 from sqlmesh.core.model.kind import TimeColumn
 from sqlmesh.core.schema_diff import SchemaDiffer
@@ -911,7 +911,7 @@ def scd_type_2(
         self,
         target_table: TableName,
         source_table: QueryOrDF,
-        unique_key: t.Sequence[str],
+        unique_key: t.Sequence[exp.Expression],
         valid_from_name: str,
         valid_to_name: str,
         updated_at_name: str,
@@ -937,7 +937,7 @@ def scd_type_2(
                 exp.Select()  # type: ignore
                 .with_(
                     "source",
-                    exp.select(*unmanaged_columns)
+                    exp.select(exp.true().as_("_exists"), *unmanaged_columns)
                     .distinct(*unique_key)
                     .from_(source_query.subquery("raw_source")),  # type: ignore
                 )
@@ -964,8 +964,8 @@ def scd_type_2(
                         "latest",
                         on=exp.and_(
                             *[
-                                exp.column(col, table="static").eq(exp.column(col, table="latest"))
-                                for col in unique_key
+                                add_table(key, "static").eq(add_table(key, "latest"))
+                                for key in unique_key
                             ]
                         ),
                         join_type="left",
@@ -976,7 +976,8 @@ def scd_type_2(
                 .with_(
                     "latest_deleted",
                     exp.select(
-                        *unique_key,
+                        exp.true().as_("_exists"),
+                        *(part.as_(f"_key{i}") for i, part in enumerate(unique_key)),
                         f"MAX({valid_to_name}) AS {valid_to_name}",
                     )
                     .from_("deleted")
@@ -987,34 +988,43 @@ def scd_type_2(
                 .with_(
                     "joined",
                     exp.select(
-                        *(f"latest.{col} AS t_{col}" for col in columns_to_types),
-                        *(f"source.{col} AS s_{col}" for col in unmanaged_columns),
+                        exp.column("_exists", table="source"),
+                        *(
+                            exp.column(col, table="latest").as_(f"t_{col}")
+                            for col in columns_to_types
+                        ),
+                        *(exp.column(col, table="source").as_(col) for col in unmanaged_columns),
                     )
                     .from_("latest")
                     .join(
                         "source",
                         on=exp.and_(
                             *[
-                                exp.column(col, table="latest").eq(exp.column(col, table="source"))
-                                for col in unique_key
+                                add_table(key, "latest").eq(add_table(key, "source"))
+                                for key in unique_key
                             ]
                         ),
                         join_type="left",
                     )
                     .union(
                         exp.select(
-                            *(f"latest.{col} AS t_{col}" for col in columns_to_types),
-                            *(f"source.{col} AS s_{col}" for col in unmanaged_columns),
+                            exp.column("_exists", table="source"),
+                            *(
+                                exp.column(col, table="latest").as_(f"t_{col}")
+                                for col in columns_to_types
+                            ),
+                            *(
+                                exp.column(col, table="source").as_(col)
+                                for col in unmanaged_columns
+                            ),
                         )
                         .from_("latest")
                         .join(
                             "source",
                             on=exp.and_(
                                 *[
-                                    exp.column(col, table="latest").eq(
-                                        exp.column(col, table="source")
-                                    )
-                                    for col in unique_key
+                                    add_table(key, "latest").eq(add_table(key, "source"))
+                                    for key in unique_key
                                 ]
                             ),
                             join_type="right",
@@ -1025,25 +1035,32 @@ def scd_type_2(
                 .with_(
                     "updated_rows",
                     exp.select(
-                        *(f"COALESCE(t_{col}, s_{col}) as {col}" for col in unmanaged_columns),
+                        *(
+                            exp.func(
+                                "COALESCE",
+                                exp.column(f"t_{col}", table="joined"),
+                                exp.column(col, table="joined"),
+                            ).as_(col)
+                            for col in unmanaged_columns
+                        ),
                         f"""
                         CASE
                             WHEN t_{valid_from_name} IS NULL
-                                 AND latest_deleted.{unique_key[0]} IS NOT NULL
+                                 AND latest_deleted._exists IS NOT NULL
                             THEN CASE
-                                    WHEN latest_deleted.{valid_to_name} > s_{updated_at_name}
+                                    WHEN latest_deleted.{valid_to_name} > {updated_at_name}
                                     THEN latest_deleted.{valid_to_name}
-                                    ELSE s_{updated_at_name}
+                                    ELSE {updated_at_name}
                                  END
                             WHEN t_{valid_from_name} IS NULL
                             THEN {self._to_utc_timestamp('1970-01-01 00:00:00+00:00')}
                             ELSE t_{valid_from_name}
                         END AS {valid_from_name}""",
                         f"""
                         CASE
-                            WHEN s_{updated_at_name} > t_{updated_at_name}
-                            THEN s_{updated_at_name}
-                            WHEN s_{unique_key[0]} IS NULL
+                            WHEN {updated_at_name} > t_{updated_at_name}
+                            THEN {updated_at_name}
+                            WHEN joined._exists IS NULL
                             THEN {self._to_utc_timestamp(to_ts(execution_time))}
                             ELSE t_{valid_to_name}
                         END AS {valid_to_name}""",
@@ -1053,10 +1070,10 @@ def scd_type_2(
                         "latest_deleted",
                         on=exp.and_(
                             *[
-                                exp.column(f"s_{col}", table="joined").eq(
-                                    exp.column(col, table="latest_deleted")
+                                add_table(part, "joined").eq(
+                                    exp.column(f"_key{i}", "latest_deleted")
                                 )
-                                for col in unique_key
+                                for i, part in enumerate(unique_key)
                             ]
                         ),
                         join_type="left",
@@ -1066,14 +1083,12 @@ def scd_type_2(
                 .with_(
                     "inserted_rows",
                     exp.select(
-                        *(f"s_{col} as {col}" for col in unmanaged_columns),
-                        f"s_{updated_at_name} as {valid_from_name}",
+                        *unmanaged_columns,
+                        f"{updated_at_name} as {valid_from_name}",
                         f"{self._to_utc_timestamp(exp.null())} as {valid_to_name}",
                     )
                     .from_("joined")
-                    .where(
-                        f"t_{unique_key[0]} IS NOT NULL AND s_{unique_key[0]} IS NOT NULL AND s_{updated_at_name} > t_{updated_at_name}"
-                    ),
+                    .where(f"{updated_at_name} > t_{updated_at_name}"),
                 )
                 .select("*")
                 .from_("static")
@@ -1097,18 +1112,15 @@ def merge(
         target_table: TableName,
         source_table: QueryOrDF,
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
-        unique_key: t.Sequence[str],
+        unique_key: t.Sequence[exp.Expression],
     ) -> None:
         source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
             source_table, columns_to_types, target_table=target_table
         )
         columns_to_types = columns_to_types or self.columns(target_table)
         on = exp.and_(
             *(
-                exp.EQ(
-                    this=exp.column(part, MERGE_TARGET_ALIAS),
-                    expression=exp.column(part, MERGE_SOURCE_ALIAS),
-                )
+                add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS))
                 for part in unique_key
             )
         )

sqlmesh/core/engine_adapter/mixins.py

@@ -23,7 +23,7 @@ def merge(
         target_table: TableName,
         source_table: QueryOrDF,
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
-        unique_key: t.Sequence[str],
+        unique_key: t.Sequence[exp.Expression],
     ) -> None:
         """
         Merge implementation for engine adapters that do not support merge natively.

sqlmesh/core/model/common.py

@@ -3,7 +3,8 @@
 import typing as t
 
 from sqlglot import exp
-from sqlglot.helper import seq_get
+from sqlglot.helper import ensure_list, seq_get
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
 from sqlmesh.core.dialect import parse, parse_one
 from sqlmesh.utils import str_to_bool
@@ -45,6 +46,32 @@ def parse_expression(
     return v
 
 
+@field_validator_v1_args
+def parse_expressions(cls: t.Type, v: t.Any, values: t.Dict[str, t.Any]) -> t.List[exp.Expression]:
+    dialect = values.get("dialect")
+
+    if isinstance(v, (exp.Tuple, exp.Array)):
+        expressions: t.List[exp.Expression] = v.expressions
+    elif isinstance(v, exp.Expression):
+        expressions = [v]
+    else:
+        expressions = [
+            parse_one(entry, dialect=dialect) if isinstance(entry, str) else entry
+            for entry in ensure_list(v)
+        ]
+
+    results = []
+
+    for expr in expressions:
+        expr = normalize_identifiers(
+            exp.to_column(expr.name) if isinstance(expr, exp.Identifier) else expr
+        )
+        expr.meta["dialect"] = dialect
+        results.append(expr)
+
+    return results
+
+
 def parse_bool(v: t.Any) -> bool:
     if isinstance(v, exp.Boolean):
         return v.this
@@ -92,6 +119,7 @@ def parse_properties(cls: t.Type, v: t.Any, values: t.Dict[str, t.Any]) -> t.Opt
     "expressions_",
     "pre_statements_",
     "post_statements_",
+    "unique_key",
     mode="before",
     check_fields=False,
 )(parse_expression)

sqlmesh/core/model/definition.py

@@ -721,7 +721,7 @@ def _data_hash_values(self) -> t.List[str]:
             data.append(self.kind.time_column.column)
             data.append(self.kind.time_column.format)
         elif isinstance(self.kind, IncrementalByUniqueKeyKind):
-            data.extend(self.kind.unique_key)
+            data.extend((k.sql() for k in self.kind.unique_key))
 
         return data  # type: ignore
 

sqlmesh/core/model/kind.py

@@ -9,13 +9,12 @@
 from sqlglot.time import format_time
 
 from sqlmesh.core import dialect as d
-from sqlmesh.core.model.common import parse_properties
+from sqlmesh.core.model.common import parse_expressions, parse_properties
 from sqlmesh.core.model.seed import CsvSettings
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import (
     PydanticModel,
     SQLGlotBool,
-    SQLGlotListOfStrings,
     SQLGlotPositiveInt,
     SQLGlotString,
     field_validator,
@@ -104,15 +103,7 @@ def model_kind_name(self) -> t.Optional[ModelKindName]:
         return self
 
 
-def _unique_key_validator(v: t.Any) -> t.List[str]:
-    if isinstance(v, exp.Identifier):
-        return [v.name]
-    if isinstance(v, (exp.Tuple, exp.Array)):
-        return [e.name for e in v.expressions]
-    return [i.name if isinstance(i, exp.Identifier) else str(i) for i in v]
-
-
-unique_key_validator = field_validator("unique_key", mode="before")(_unique_key_validator)
+_unique_key_validator = field_validator("unique_key", mode="before")(parse_expressions)
 
 
 class _ModelKind(PydanticModel, ModelKindMixin):
@@ -216,7 +207,8 @@ def to_expression(self, dialect: str = "", **kwargs: t.Any) -> d.ModelKind:
 
 class IncrementalByUniqueKeyKind(_Incremental):
     name: Literal[ModelKindName.INCREMENTAL_BY_UNIQUE_KEY] = ModelKindName.INCREMENTAL_BY_UNIQUE_KEY
-    unique_key: SQLGlotListOfStrings
+    unique_key: t.List[exp.Expression]
+    _unique_key_validator = _unique_key_validator
 
 
 class IncrementalUnmanagedKind(_ModelKind):
@@ -281,13 +273,14 @@ class FullKind(_ModelKind):
 
 class SCDType2Kind(_ModelKind):
     name: Literal[ModelKindName.SCD_TYPE_2] = ModelKindName.SCD_TYPE_2
-    unique_key: SQLGlotListOfStrings
+    unique_key: t.List[exp.Expression]
     valid_from_name: SQLGlotString = "valid_from"
     valid_to_name: SQLGlotString = "valid_to"
     updated_at_name: SQLGlotString = "updated_at"
 
     forward_only: SQLGlotBool = True
     disable_restatement: SQLGlotBool = True
+    _unique_key_validator = _unique_key_validator
 
     @property
     def managed_columns(self) -> t.Dict[str, exp.DataType]:

sqlmesh/core/model/meta.py

@@ -8,7 +8,11 @@
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
 from sqlmesh.core import dialect as d
-from sqlmesh.core.model.common import bool_validator, properties_validator
+from sqlmesh.core.model.common import (
+    bool_validator,
+    parse_expressions,
+    properties_validator,
+)
 from sqlmesh.core.model.kind import (
     IncrementalByUniqueKeyKind,
     ModelKind,
@@ -139,27 +143,9 @@ def _string_validator(cls, v: t.Any) -> t.Optional[str]:
     def _partition_by_validator(
         cls, v: t.Any, values: t.Dict[str, t.Any]
     ) -> t.List[exp.Expression]:
-        dialect = values.get("dialect")
-
-        if isinstance(v, (exp.Tuple, exp.Array)):
-            partitions: t.List[exp.Expression] = v.expressions
-        elif isinstance(v, exp.Expression):
-            partitions = [v]
-        else:
-            partitions = [
-                d.parse_one(entry, dialect=dialect) if isinstance(entry, str) else entry
-                for entry in ensure_list(v)
-            ]
-
-        partitions = [
-            normalize_identifiers(
-                exp.to_column(expr.name) if isinstance(expr, exp.Identifier) else expr
-            )
-            for expr in partitions
-        ]
+        partitions = parse_expressions(cls, v, values)
 
         for partition in partitions:
-            partition.meta["dialect"] = dialect
             num_cols = len(list(partition.find_all(exp.Column)))
 
             error_msg: t.Optional[str] = None
@@ -289,7 +275,7 @@ def time_column(self) -> t.Optional[TimeColumn]:
         return getattr(self.kind, "time_column", None)
 
     @property
-    def unique_key(self) -> t.List[str]:
+    def unique_key(self) -> t.List[exp.Expression]:
         if isinstance(self.kind, (IncrementalByUniqueKeyKind, SCDType2Kind)):
             return self.kind.unique_key
         return []