Repository: TobikoData/sqlmesh
Pull request #3124 (2024-09-12 20:02:11)

feat: support multiple when matched inc unique key
Prior to this PR the when matched expression only supported a single expression. This PR makes it so you can pass in multiple. We assume the underlying engine will support it (will get an error at runtime if that is not the case). 


Patch:
diff --git a/docs/concepts/models/model_kinds.md b/docs/concepts/models/model_kinds.md
index 2b9047e3..198daf92 100644
--- a/docs/concepts/models/model_kinds.md
+++ b/docs/concepts/models/model_kinds.md
@@ -259,6 +259,19 @@ MODEL (
 
 The `source` and `target` aliases are required when using the `when_matched` expression in order to distinguish between the source and target columns.
 
+Multiple `WHEN MATCHED` expressions can also be provided. Ex:
+
+```sql linenums="1" hl_lines="5-6"
+MODEL (
+  name db.employees,
+  kind INCREMENTAL_BY_UNIQUE_KEY (
+    unique_key name,
+    when_matched WHEN MATCHED AND source.value IS NULL THEN UPDATE SET target.salary = COALESCE(source.salary, target.salary),
+    WHEN MATCHED THEN UPDATE SET target.title = COALESCE(source.title, target.title)
+  )
+);
+```
+
 **Note**: `when_matched` is only available on engines that support the `MERGE` statement. Currently supported engines include:
 
 * BigQuery
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 164c87ac..1d60878b 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -398,7 +398,9 @@ def _parse_props(self: Parser) -> t.Optional[exp.Expression]:
 
     name = key.name.lower()
     if name == "when_matched":
-        value: t.Optional[exp.Expression] = self._parse_when_matched()[0]
+        value: t.Optional[t.Union[exp.Expression, t.List[exp.Expression]]] = (
+            self._parse_when_matched()  # type: ignore
+        )
     elif name == "time_data_type":
         # TODO: if we make *_data_type a convention to parse things into exp.DataType, we could make this more generic
         value = self._parse_types(schema=True)
@@ -410,7 +412,7 @@ def _parse_props(self: Parser) -> t.Optional[exp.Expression]:
 
     if name == "path" and value:
         # Make sure if we get a windows path that it is converted to posix
-        value = exp.Literal.string(value.this.replace("\\", "/"))
+        value = exp.Literal.string(value.this.replace("\\", "/"))  # type: ignore
 
     return self.expression(exp.Property, this=name, value=value)
 
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index f16b8331..c978fa8d 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -1726,7 +1726,7 @@ class EngineAdapter:
         source_table: QueryOrDF,
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
         unique_key: t.Sequence[exp.Expression],
-        when_matched: t.Optional[exp.When] = None,
+        when_matched: t.Optional[t.Union[exp.When, t.List[exp.When]]] = None,
     ) -> None:
         source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
             source_table, columns_to_types, target_table=target_table
@@ -1749,6 +1749,7 @@ class EngineAdapter:
                     ],
                 ),
             )
+        when_matched = ensure_list(when_matched)
         when_not_matched = exp.When(
             matched=False,
             source=False,
@@ -1759,13 +1760,14 @@ class EngineAdapter:
                 ),
             ),
         )
+        match_expressions = when_matched + [when_not_matched]
         for source_query in source_queries:
             with source_query as query:
                 self._merge(
                     target_table=target_table,
                     query=query,
                     on=on,
-                    match_expressions=[when_matched, when_not_matched],
+                    match_expressions=match_expressions,
                 )
 
     def rename_table(
diff --git a/sqlmesh/core/engine_adapter/mixins.py b/sqlmesh/core/engine_adapter/mixins.py
index da418506..90e274dc 100644
--- a/sqlmesh/core/engine_adapter/mixins.py
+++ b/sqlmesh/core/engine_adapter/mixins.py
@@ -25,7 +25,7 @@ class LogicalMergeMixin(EngineAdapter):
         source_table: QueryOrDF,
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
         unique_key: t.Sequence[exp.Expression],
-        when_matched: t.Optional[exp.When] = None,
+        when_matched: t.Optional[t.Union[exp.When, t.List[exp.When]]] = None,
     ) -> None:
         """
         Merge implementation for engine adapters that do not support merge natively.
diff --git a/sqlmesh/core/model/kind.py b/sqlmesh/core/model/kind.py
index d7a6e504..06e16dc7 100644
--- a/sqlmesh/core/model/kind.py
+++ b/sqlmesh/core/model/kind.py
@@ -6,6 +6,7 @@ from enum import Enum
 
 from pydantic import Field
 from sqlglot import exp
+from sqlglot.helper import ensure_list
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 from sqlglot.optimizer.simplify import gen
@@ -424,14 +425,16 @@ class IncrementalByTimeRangeKind(_IncrementalBy):
 class IncrementalByUniqueKeyKind(_IncrementalBy):
     name: Literal[ModelKindName.INCREMENTAL_BY_UNIQUE_KEY] = ModelKindName.INCREMENTAL_BY_UNIQUE_KEY
     unique_key: SQLGlotListOfFields
-    when_matched: t.Optional[exp.When] = None
+    when_matched: t.Optional[t.List[exp.When]] = None
     batch_concurrency: Literal[1] = 1
 
     @field_validator("when_matched", mode="before")
     @field_validator_v1_args
     def _when_matched_validator(
-        cls, v: t.Optional[t.Union[exp.When, str]], values: t.Dict[str, t.Any]
-    ) -> t.Optional[exp.When]:
+        cls,
+        v: t.Optional[t.Union[exp.When, str, t.List[exp.When], t.List[str]]],
+        values: t.Dict[str, t.Any],
+    ) -> t.Optional[t.List[exp.When]]:
         def replace_table_references(expression: exp.Expression) -> exp.Expression:
             from sqlmesh.core.engine_adapter.base import (
                 MERGE_SOURCE_ALIAS,
@@ -451,13 +454,19 @@ class IncrementalByUniqueKeyKind(_IncrementalBy):
                     )
             return expression
 
-        if isinstance(v, str):
-            return t.cast(exp.When, d.parse_one(v, into=exp.When, dialect=get_dialect(values)))
-
         if not v:
-            return v
-
-        return t.cast(exp.When, v.transform(replace_table_references))
+            return v  # type: ignore
+
+        result = []
+        list_v = ensure_list(v)
+        for value in ensure_list(list_v):
+            if isinstance(value, str):
+                result.append(
+                    t.cast(exp.When, d.parse_one(value, into=exp.When, dialect=get_dialect(values)))
+                )
+            else:
+                result.append(t.cast(exp.When, value.transform(replace_table_references)))  # type: ignore
+        return result
 
     @property
     def data_hash_values(self) -> t.List[t.Optional[str]]:
diff --git a/sqlmesh/core/model/meta.py b/sqlmesh/core/model/meta.py
index 3995bd2b..2c6a0b67 100644
--- a/sqlmesh/core/model/meta.py
+++ b/sqlmesh/core/model/meta.py
@@ -439,7 +439,7 @@ class ModelMeta(_Node):
         return getattr(self.kind, "managed_columns", {})
 
     @property
-    def when_matched(self) -> t.Optional[exp.When]:
+    def when_matched(self) -> t.Optional[t.List[exp.When]]:
         if isinstance(self.kind, IncrementalByUniqueKeyKind):
             return self.kind.when_matched
         return None


Test patch:
diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py
index 863e5793..4785ecfc 100644
--- a/tests/core/engine_adapter/test_base.py
+++ b/tests/core/engine_adapter/test_base.py
@@ -974,6 +974,75 @@ MERGE INTO "target" AS "__MERGE_TARGET__" USING (
     )
 
 
+def test_merge_when_matched_multiple(make_mocked_engine_adapter: t.Callable, assert_exp_eq):
+    adapter = make_mocked_engine_adapter(EngineAdapter)
+
+    adapter.merge(
+        target_table="target",
+        source_table=t.cast(exp.Select, parse_one('SELECT "ID", ts, val FROM source')),
+        columns_to_types={
+            "ID": exp.DataType.build("int"),
+            "ts": exp.DataType.build("timestamp"),
+            "val": exp.DataType.build("int"),
+        },
+        unique_key=[exp.to_identifier("ID", quoted=True)],
+        when_matched=[
+            exp.When(
+                matched=True,
+                condition=exp.column("ID", "__MERGE_SOURCE__").eq(exp.Literal.number(1)),
+                then=exp.Update(
+                    expressions=[
+                        exp.column("val", "__MERGE_TARGET__").eq(
+                            exp.column("val", "__MERGE_SOURCE__")
+                        ),
+                        exp.column("ts", "__MERGE_TARGET__").eq(
+                            exp.Coalesce(
+                                this=exp.column("ts", "__MERGE_SOURCE__"),
+                                expressions=[exp.column("ts", "__MERGE_TARGET__")],
+                            )
+                        ),
+                    ],
+                ),
+            ),
+            exp.When(
+                matched=True,
+                source=False,
+                then=exp.Update(
+                    expressions=[
+                        exp.column("val", "__MERGE_TARGET__").eq(
+                            exp.column("val", "__MERGE_SOURCE__")
+                        ),
+                        exp.column("ts", "__MERGE_TARGET__").eq(
+                            exp.Coalesce(
+                                this=exp.column("ts", "__MERGE_SOURCE__"),
+                                expressions=[exp.column("ts", "__MERGE_TARGET__")],
+                            )
+                        ),
+                    ],
+                ),
+            ),
+        ],
+    )
+
+    assert_exp_eq(
+        adapter.cursor.execute.call_args[0][0],
+        """
+MERGE INTO "target" AS "__MERGE_TARGET__" USING (
+  SELECT
+    "ID",
+    "ts",
+    "val"
+  FROM "source"
+) AS "__MERGE_SOURCE__"
+  ON "__MERGE_TARGET__"."ID" = "__MERGE_SOURCE__"."ID"
+  WHEN MATCHED AND "__MERGE_SOURCE__"."ID" = 1 THEN UPDATE SET "__MERGE_TARGET__"."val" = "__MERGE_SOURCE__"."val", "__MERGE_TARGET__"."ts" = COALESCE("__MERGE_SOURCE__"."ts", "__MERGE_TARGET__"."ts"),
+  WHEN MATCHED THEN UPDATE SET "__MERGE_TARGET__"."val" = "__MERGE_SOURCE__"."val", "__MERGE_TARGET__"."ts" = COALESCE("__MERGE_SOURCE__"."ts", "__MERGE_TARGET__"."ts")
+  WHEN NOT MATCHED THEN INSERT ("ID", "ts", "val")
+    VALUES ("__MERGE_SOURCE__"."ID", "__MERGE_SOURCE__"."ts", "__MERGE_SOURCE__"."val")
+""",
+    )
+
+
 def test_scd_type_2_by_time(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(EngineAdapter)
 
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index de953650..51217d1a 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -3904,10 +3904,44 @@ def test_when_matched():
     expected_when_matched = "WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.salary = COALESCE(__MERGE_SOURCE__.salary, __MERGE_TARGET__.salary)"
 
     model = load_sql_based_model(expressions, dialect="hive")
-    assert model.kind.when_matched.sql() == expected_when_matched
+    assert len(model.kind.when_matched) == 1
+    assert model.kind.when_matched[0].sql() == expected_when_matched
 
     model = SqlModel.parse_raw(model.json())
-    assert model.kind.when_matched.sql() == expected_when_matched
+    assert len(model.kind.when_matched) == 1
+    assert model.kind.when_matched[0].sql() == expected_when_matched
+
+
+def test_when_matched_multiple():
+    expressions = d.parse(
+        """
+        MODEL (
+          name db.employees,
+          kind INCREMENTAL_BY_UNIQUE_KEY (
+            unique_key name,
+            when_matched WHEN MATCHED AND source.x = 1 THEN UPDATE SET target.salary = COALESCE(source.salary, target.salary),
+            WHEN MATCHED THEN UPDATE SET target.salary = COALESCE(source.salary, target.salary)
+            
+          )
+        );
+        SELECT 'name' AS name, 1 AS salary;
+    """
+    )
+
+    expected_when_matched = [
+        "WHEN MATCHED AND __MERGE_SOURCE__.x = 1 THEN UPDATE SET __MERGE_TARGET__.salary = COALESCE(__MERGE_SOURCE__.salary, __MERGE_TARGET__.salary)",
+        "WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.salary = COALESCE(__MERGE_SOURCE__.salary, __MERGE_TARGET__.salary)",
+    ]
+
+    model = load_sql_based_model(expressions, dialect="hive")
+    assert len(model.kind.when_matched) == 2
+    assert model.kind.when_matched[0].sql() == expected_when_matched[0]
+    assert model.kind.when_matched[1].sql() == expected_when_matched[1]
+
+    model = SqlModel.parse_raw(model.json())
+    assert len(model.kind.when_matched) == 2
+    assert model.kind.when_matched[0].sql() == expected_when_matched[0]
+    assert model.kind.when_matched[1].sql() == expected_when_matched[1]
 
 
 def test_default_catalog_sql(assert_exp_eq):
@@ -5438,7 +5472,35 @@ on_destructive_change 'ERROR'
         .sql()
         == """INCREMENTAL_BY_UNIQUE_KEY (
 unique_key ("a"),
-when_matched WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.b = COALESCE(__MERGE_SOURCE__.b, __MERGE_TARGET__.b),
+when_matched ARRAY(WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.b = COALESCE(__MERGE_SOURCE__.b, __MERGE_TARGET__.b)),
+batch_concurrency 1,
+forward_only FALSE,
+disable_restatement FALSE,
+on_destructive_change 'ERROR'
+)"""
+    )
+
+    assert (
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name db.table,
+                kind INCREMENTAL_BY_UNIQUE_KEY(
+                    unique_key a,
+                    when_matched WHEN MATCHED AND source.x = 1 THEN UPDATE SET target.b = COALESCE(source.b, target.b),
+                    WHEN MATCHED THEN UPDATE SET target.b = COALESCE(source.b, target.b)
+                ),
+            );
+            SELECT a, b
+            """
+            )
+        )
+        .kind.to_expression()
+        .sql()
+        == """INCREMENTAL_BY_UNIQUE_KEY (
+unique_key ("a"),
+when_matched ARRAY(WHEN MATCHED AND __MERGE_SOURCE__.x = 1 THEN UPDATE SET __MERGE_TARGET__.b = COALESCE(__MERGE_SOURCE__.b, __MERGE_TARGET__.b), WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.b = COALESCE(__MERGE_SOURCE__.b, __MERGE_TARGET__.b)),
 batch_concurrency 1,
 forward_only FALSE,
 disable_restatement FALSE,
diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
index c76151a4..12009a4d 100644
--- a/tests/core/test_snapshot_evaluator.py
+++ b/tests/core/test_snapshot_evaluator.py
@@ -1821,23 +1821,104 @@ def test_create_incremental_by_unique_key_updated_at_exp(adapter_mock, make_snap
             "updated_at": exp.DataType.build("TIMESTAMP"),
         },
         unique_key=[exp.to_column("id", quoted=True)],
-        when_matched=exp.When(
-            matched=True,
-            source=False,
-            then=exp.Update(
-                expressions=[
-                    exp.column("name", MERGE_TARGET_ALIAS).eq(
-                        exp.column("name", MERGE_SOURCE_ALIAS)
-                    ),
-                    exp.column("updated_at", MERGE_TARGET_ALIAS).eq(
-                        exp.Coalesce(
-                            this=exp.column("updated_at", MERGE_SOURCE_ALIAS),
-                            expressions=[exp.column("updated_at", MERGE_TARGET_ALIAS)],
-                        )
-                    ),
-                ],
+        when_matched=[
+            exp.When(
+                matched=True,
+                source=False,
+                then=exp.Update(
+                    expressions=[
+                        exp.column("name", MERGE_TARGET_ALIAS).eq(
+                            exp.column("name", MERGE_SOURCE_ALIAS)
+                        ),
+                        exp.column("updated_at", MERGE_TARGET_ALIAS).eq(
+                            exp.Coalesce(
+                                this=exp.column("updated_at", MERGE_SOURCE_ALIAS),
+                                expressions=[exp.column("updated_at", MERGE_TARGET_ALIAS)],
+                            )
+                        ),
+                    ],
+                ),
+            )
+        ],
+    )
+
+
+def test_create_incremental_by_unique_key_multiple_updated_at_exp(adapter_mock, make_snapshot):
+    evaluator = SnapshotEvaluator(adapter_mock)
+    model = load_sql_based_model(
+        parse(  # type: ignore
+            """
+            MODEL (
+                name test_schema.test_model,
+                kind INCREMENTAL_BY_UNIQUE_KEY (
+                    unique_key [id],
+                    when_matched WHEN MATCHED AND source.id = 1 THEN UPDATE SET target.name = source.name, target.updated_at = COALESCE(source.updated_at, target.updated_at),
+                    WHEN MATCHED THEN UPDATE SET target.name = source.name, target.updated_at = COALESCE(source.updated_at, target.updated_at)
+                )
+            );
+
+            SELECT id::int, name::string, updated_at::timestamp FROM tbl;
+            """
+        )
+    )
+
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot.intervals = [(to_timestamp("2020-01-01"), to_timestamp("2020-01-02"))]
+
+    evaluator.evaluate(
+        snapshot,
+        start="2020-01-01",
+        end="2020-01-02",
+        execution_time="2020-01-02",
+        snapshots={},
+    )
+
+    adapter_mock.merge.assert_called_once_with(
+        snapshot.table_name(),
+        model.render_query(),
+        columns_to_types={
+            "id": exp.DataType.build("INT"),
+            "name": exp.DataType.build("STRING"),
+            "updated_at": exp.DataType.build("TIMESTAMP"),
+        },
+        unique_key=[exp.to_column("id", quoted=True)],
+        when_matched=[
+            exp.When(
+                matched=True,
+                condition=exp.column("id", MERGE_SOURCE_ALIAS).eq(exp.Literal.number(1)),
+                then=exp.Update(
+                    expressions=[
+                        exp.column("name", MERGE_TARGET_ALIAS).eq(
+                            exp.column("name", MERGE_SOURCE_ALIAS)
+                        ),
+                        exp.column("updated_at", MERGE_TARGET_ALIAS).eq(
+                            exp.Coalesce(
+                                this=exp.column("updated_at", MERGE_SOURCE_ALIAS),
+                                expressions=[exp.column("updated_at", MERGE_TARGET_ALIAS)],
+                            )
+                        ),
+                    ],
+                ),
             ),
-        ),
+            exp.When(
+                matched=True,
+                source=False,
+                then=exp.Update(
+                    expressions=[
+                        exp.column("name", MERGE_TARGET_ALIAS).eq(
+                            exp.column("name", MERGE_SOURCE_ALIAS)
+                        ),
+                        exp.column("updated_at", MERGE_TARGET_ALIAS).eq(
+                            exp.Coalesce(
+                                this=exp.column("updated_at", MERGE_SOURCE_ALIAS),
+                                expressions=[exp.column("updated_at", MERGE_TARGET_ALIAS)],
+                            )
+                        ),
+                    ],
+                ),
+            ),
+        ],
     )