Fix: Deserialization of the when_matched attribute of the INCREMENTAL_BY_UNIQUE_KEY kind (#1851)

Author: izeigerman

sqlmesh/core/dialect.py

@@ -11,6 +11,7 @@
 from sqlglot import Dialect, Generator, ParseError, Parser, Tokenizer, TokenType, exp
 from sqlglot.dialects.dialect import DialectType
 from sqlglot.dialects.snowflake import Snowflake
+from sqlglot.helper import seq_get
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.scope import traverse_scope
 from sqlglot.tokens import Token
@@ -607,8 +608,10 @@ class ChunkType(Enum):
     SQL = auto()
 
 
-def parse_one(sql: str, dialect: t.Optional[str] = None) -> exp.Expression:
-    expressions = parse(sql, default_dialect=dialect, match_dialect=False)
+def parse_one(
+    sql: str, dialect: t.Optional[str] = None, into: t.Optional[exp.IntoType] = None
+) -> exp.Expression:
+    expressions = parse(sql, default_dialect=dialect, match_dialect=False, into=into)
     if not expressions:
         raise SQLMeshError(f"No expressions found in '{sql}'")
     elif len(expressions) > 1:
@@ -617,7 +620,10 @@ def parse_one(sql: str, dialect: t.Optional[str] = None) -> exp.Expression:
 
 
 def parse(
-    sql: str, default_dialect: t.Optional[str] = None, match_dialect: bool = True
+    sql: str,
+    default_dialect: t.Optional[str] = None,
+    match_dialect: bool = True,
+    into: t.Optional[exp.IntoType] = None,
 ) -> t.List[exp.Expression]:
     """Parse a sql string.
 
@@ -668,7 +674,10 @@ def parse(
 
     for chunk, chunk_type in chunks:
         if chunk_type == ChunkType.SQL:
-            for expression in parser.parse(chunk, sql):
+            parsed_expressions: t.List[t.Optional[exp.Expression]] = (
+                parser.parse(chunk, sql) if into is None else parser.parse_into(into, chunk, sql)
+            )
+            for expression in parsed_expressions:
                 if expression:
                     expression.meta["sql"] = parser._find_sql(chunk[0], chunk[-1])
                     expressions.append(expression)
@@ -706,6 +715,10 @@ def extend_sqlglot() -> None:
         parser.QUERY_MODIFIER_PARSERS.update(
             {TokenType.PARAMETER: lambda self: _parse_body_macro(self)}
         )
+        # FIXME: Delete the extension below after upgrading to SQLGlot >= 20.3.0.
+        parser.EXPRESSION_PARSERS.update(
+            {exp.When: lambda self: seq_get(self._parse_when_matched(), 0)}
+        )
 
     for generator in generators:
         if MacroFunc not in generator.TRANSFORMS:

sqlmesh/core/model/kind.py

@@ -231,7 +231,7 @@ class IncrementalByUniqueKeyKind(_Incremental):
     @field_validator("when_matched", mode="before")
     @field_validator_v1_args
     def _when_matched_validator(
-        cls, v: t.Optional[exp.When], values: t.Dict[str, t.Any]
+        cls, v: t.Optional[t.Union[exp.When, str]], values: t.Dict[str, t.Any]
     ) -> t.Optional[exp.When]:
         def replace_table_references(expression: exp.Expression) -> exp.Expression:
             from sqlmesh.core.engine_adapter.base import (
@@ -252,9 +252,12 @@ def replace_table_references(expression: exp.Expression) -> exp.Expression:
                     )
             return expression
 
+        if isinstance(v, str):
+            return t.cast(exp.When, d.parse_one(v, into=exp.When))
+
         if not v:
             return v
-        v.meta["dialect"] = values.get("dialect")
+
         return v.transform(replace_table_references)
 
 

tests/core/test_model.py

@@ -2556,3 +2556,26 @@ def test_null_column_type():
         "id": exp.DataType.build("int"),
     }
     assert not model.annotated
+
+
+def test_when_matched():
+    expressions = d.parse(
+        """
+        MODEL (
+          name db.employees,
+          kind INCREMENTAL_BY_UNIQUE_KEY (
+            unique_key name,
+            when_matched WHEN MATCHED THEN UPDATE SET target.salary = COALESCE(source.salary, target.salary)
+          )
+        );
+        SELECT 'name' AS name, 1 AS salary;
+    """
+    )
+
+    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
+
+    model = SqlModel.parse_raw(model.json())
+    assert model.kind.when_matched.sql() == expected_when_matched