Fix: Normalize when_matched and merge_filter expressions to the source dialect

Author: erindru

sqlmesh/core/dialect.py

@@ -1388,17 +1388,26 @@ def is_meta_expression(v: t.Any) -> bool:
     return isinstance(v, (Audit, Metric, Model))
 
 
-def replace_merge_table_aliases(expression: exp.Expression) -> exp.Expression:
+def replace_merge_table_aliases(
+    expression: exp.Expression, dialect: t.Optional[str] = None
+) -> exp.Expression:
     """
     Resolves references from the "source" and "target" tables (or their DBT equivalents)
     with the corresponding SQLMesh merge aliases (MERGE_SOURCE_ALIAS and MERGE_TARGET_ALIAS)
     """
     from sqlmesh.core.engine_adapter.base import MERGE_SOURCE_ALIAS, MERGE_TARGET_ALIAS
 
+    normalized_merge_source_alias = normalize_identifiers(
+        exp.to_identifier(MERGE_SOURCE_ALIAS), dialect
+    )
+    normalized_merge_target_alias = normalize_identifiers(
+        exp.to_identifier(MERGE_TARGET_ALIAS), dialect
+    )
+
     if isinstance(expression, exp.Column) and (first_part := expression.parts[0]):
         if first_part.this.lower() in ("target", "dbt_internal_dest", "__merge_target__"):
-            first_part.replace(exp.to_identifier(MERGE_TARGET_ALIAS))
+            first_part.replace(normalized_merge_target_alias)
         elif first_part.this.lower() in ("source", "dbt_internal_source", "__merge_source__"):
-            first_part.replace(exp.to_identifier(MERGE_SOURCE_ALIAS))
+            first_part.replace(normalized_merge_source_alias)
 
     return expression

sqlmesh/core/model/definition.py

@@ -653,7 +653,7 @@ def render_merge_filter(
         )
         if len(rendered_exprs) != 1:
             raise SQLMeshError(f"Expected one expression but got {len(rendered_exprs)}")
-        return rendered_exprs[0].transform(d.replace_merge_table_aliases)
+        return rendered_exprs[0].transform(d.replace_merge_table_aliases, dialect=self.dialect)
 
     def _render_properties(
         self, properties: t.Dict[str, exp.Expression] | SessionProperties, **render_kwargs: t.Any

sqlmesh/core/model/kind.py

@@ -468,15 +468,20 @@ def _when_matched_validator(
             return v
         if isinstance(v, list):
             v = " ".join(v)
+
+        dialect = get_dialect(info.data)
+
         if isinstance(v, str):
             # Whens wrap the WHEN clauses, but the parentheses aren't parsed by sqlglot
             v = v.strip()
             if v.startswith("("):
                 v = v[1:-1]
 
-            return t.cast(exp.Whens, d.parse_one(v, into=exp.Whens, dialect=get_dialect(info.data)))
+            v = t.cast(exp.Whens, d.parse_one(v, into=exp.Whens, dialect=dialect))
+        else:
+            v = t.cast(exp.Whens, v.transform(d.replace_merge_table_aliases, dialect=dialect))
 
-        return t.cast(exp.Whens, v.transform(d.replace_merge_table_aliases))
+        return normalize_identifiers(v, dialect=dialect)
 
     @field_validator("merge_filter", mode="before")
     def _merge_filter_validator(
@@ -486,11 +491,16 @@ def _merge_filter_validator(
     ) -> t.Optional[exp.Expression]:
         if v is None:
             return v
+
+        dialect = get_dialect(info.data)
+
         if isinstance(v, str):
             v = v.strip()
-            return d.parse_one(v, dialect=get_dialect(info.data))
+            v = d.parse_one(v, dialect=dialect)
+        else:
+            v = v.transform(d.replace_merge_table_aliases, dialect=dialect)
 
-        return v.transform(d.replace_merge_table_aliases)
+        return normalize_identifiers(v, dialect=dialect)
 
     @property
     def data_hash_values(self) -> t.List[t.Optional[str]]:

tests/core/test_model.py

@@ -5366,7 +5366,7 @@ def test_when_matched():
     """
     )
 
-    expected_when_matched = "(WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.salary = COALESCE(__MERGE_SOURCE__.salary, __MERGE_TARGET__.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
@@ -5400,9 +5400,9 @@ def test_when_matched():
   kind INCREMENTAL_BY_UNIQUE_KEY (
     unique_key ("purchase_order_id"),
     when_matched (
-      WHEN MATCHED AND __MERGE_SOURCE__._operation = 1 THEN DELETE
-      WHEN MATCHED AND __MERGE_SOURCE__._operation <> 1 THEN UPDATE SET
-        __MERGE_TARGET__.purchase_order_id = 1
+      WHEN MATCHED AND __merge_source__._operation = 1 THEN DELETE
+      WHEN MATCHED AND __merge_source__._operation <> 1 THEN UPDATE SET
+        __merge_target__.purchase_order_id = 1
     ),
     batch_concurrency 1,
     forward_only FALSE,
@@ -5453,7 +5453,7 @@ def fingerprint_merge(
   kind INCREMENTAL_BY_UNIQUE_KEY (
     unique_key ("purchase_order_id"),
     when_matched (
-      WHEN MATCHED AND __MERGE_SOURCE__.salary <> __MERGE_TARGET__.salary THEN UPDATE SET
+      WHEN MATCHED AND __merge_source__.salary <> __merge_target__.salary THEN UPDATE SET
         ARRAY('target.update_datetime = source.update_datetime', 'target.salary = source.salary')
     ),
     batch_concurrency 1,
@@ -5487,8 +5487,8 @@ def test_when_matched_multiple():
     )
 
     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)",
+        "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", variables={"schema": "db"})
@@ -5529,13 +5529,13 @@ def test_when_matched_merge_filter_multi_part_columns():
     )
 
     expected_when_matched = [
-        "WHEN MATCHED AND __MERGE_SOURCE__.record.nested_record.field = 1 THEN UPDATE SET __MERGE_TARGET__.repeated_record.sub_repeated_record.sub_field = COALESCE(__MERGE_SOURCE__.repeated_record.sub_repeated_record.sub_field, __MERGE_TARGET__.repeated_record.sub_repeated_record.sub_field)",
-        "WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.repeated_record.sub_repeated_record.sub_field = COALESCE(__MERGE_SOURCE__.repeated_record.sub_repeated_record.sub_field, __MERGE_TARGET__.repeated_record.sub_repeated_record.sub_field)",
+        "WHEN MATCHED AND __merge_source__.record.nested_record.field = 1 THEN UPDATE SET __merge_target__.repeated_record.sub_repeated_record.sub_field = COALESCE(__merge_source__.repeated_record.sub_repeated_record.sub_field, __merge_target__.repeated_record.sub_repeated_record.sub_field)",
+        "WHEN MATCHED THEN UPDATE SET __merge_target__.repeated_record.sub_repeated_record.sub_field = COALESCE(__merge_source__.repeated_record.sub_repeated_record.sub_field, __merge_target__.repeated_record.sub_repeated_record.sub_field)",
     ]
 
     expected_merge_filter = (
-        "__MERGE_SOURCE__.record.nested_record.field < __MERGE_TARGET__.record.nested_record.field AND "
-        "__MERGE_TARGET__.repeated_record.sub_repeated_record.sub_field > __MERGE_SOURCE__.repeated_record.sub_repeated_record.sub_field"
+        "__merge_source__.record.nested_record.field < __merge_target__.record.nested_record.field AND "
+        "__merge_target__.repeated_record.sub_repeated_record.sub_field > __merge_source__.repeated_record.sub_repeated_record.sub_field"
     )
 
     model = load_sql_based_model(expressions, dialect="bigquery", variables={"schema": "db"})
@@ -5553,6 +5553,64 @@ def test_when_matched_merge_filter_multi_part_columns():
     assert model.merge_filter.sql() == expected_merge_filter
 
 
+def test_when_matched_normalization() -> None:
+    # unquoted should be normalized
+    expressions = d.parse(
+        """
+        MODEL (
+          name test.employees,
+          kind INCREMENTAL_BY_UNIQUE_KEY (
+            unique_key name,
+            when_matched (
+                WHEN MATCHED THEN UPDATE SET
+                    target.key_a = source.key_a,
+                    target.key_b = source.key_b,            
+            )
+          )
+        );
+        SELECT 'name' AS name, 1 AS key_a, 2 AS key_b;
+    """
+    )
+    model = load_sql_based_model(expressions, dialect="snowflake")
+
+    assert isinstance(model.kind, IncrementalByUniqueKeyKind)
+    assert isinstance(model.kind.when_matched, exp.Whens)
+    first_expression = model.kind.when_matched.expressions[0]
+    assert isinstance(first_expression, exp.Expression)
+    assert (
+        first_expression.sql(dialect="snowflake", identify=True)
+        == 'WHEN MATCHED THEN UPDATE SET "__MERGE_TARGET__"."KEY_A" = "__MERGE_SOURCE__"."KEY_A", "__MERGE_TARGET__"."KEY_B" = "__MERGE_SOURCE__"."KEY_B"'
+    )
+
+    # quoted should be preserved
+    expressions = d.parse(
+        """
+        MODEL (
+          name test.employees,
+          kind INCREMENTAL_BY_UNIQUE_KEY (
+            unique_key name,
+            when_matched (
+                WHEN MATCHED THEN UPDATE SET
+                    target."kEy_A" = source."kEy_A",
+                    target."kEY_b" = source.key_b,
+            )
+          )
+        );
+        SELECT 'name' AS name, 1 AS "kEy_A", 2 AS "kEY_b";
+    """
+    )
+    model = load_sql_based_model(expressions, dialect="snowflake")
+
+    assert isinstance(model.kind, IncrementalByUniqueKeyKind)
+    assert isinstance(model.kind.when_matched, exp.Whens)
+    first_expression = model.kind.when_matched.expressions[0]
+    assert isinstance(first_expression, exp.Expression)
+    assert (
+        first_expression.sql(dialect="snowflake", identify=True)
+        == 'WHEN MATCHED THEN UPDATE SET "__MERGE_TARGET__"."kEy_A" = "__MERGE_SOURCE__"."kEy_A", "__MERGE_TARGET__"."kEY_b" = "__MERGE_SOURCE__"."KEY_B"'
+    )
+
+
 def test_default_catalog_sql(assert_exp_eq):
     """
     This test validates the hashing behavior of the system as it relates to the default catalog.
@@ -7583,7 +7641,7 @@ def test_model_kind_to_expression():
         .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 (WHEN MATCHED THEN UPDATE SET __merge_target__.b = COALESCE(__merge_source__.b, __merge_target__.b)),
 batch_concurrency 1,
 forward_only FALSE,
 disable_restatement FALSE,
@@ -7611,7 +7669,7 @@ def test_model_kind_to_expression():
         .sql()
         == """INCREMENTAL_BY_UNIQUE_KEY (
 unique_key ("a"),
-when_matched (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)),
+when_matched (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,
@@ -7872,7 +7930,7 @@ def test_merge_filter():
     """
     )
 
-    expected_incremental_predicate = f"{MERGE_SOURCE_ALIAS}.salary > 0"
+    expected_incremental_predicate = f"{MERGE_SOURCE_ALIAS.lower()}.salary > 0"
 
     model = load_sql_based_model(expressions, dialect="hive")
     assert model.kind.merge_filter.sql() == expected_incremental_predicate
@@ -7913,19 +7971,19 @@ def test_merge_filter():
   kind INCREMENTAL_BY_UNIQUE_KEY (
     unique_key ("purchase_order_id"),
     when_matched (
-      WHEN MATCHED AND {MERGE_SOURCE_ALIAS}._operation = 1 THEN DELETE
-      WHEN MATCHED AND {MERGE_SOURCE_ALIAS}._operation <> 1 THEN UPDATE SET
-        {MERGE_TARGET_ALIAS}.purchase_order_id = 1
+      WHEN MATCHED AND {MERGE_SOURCE_ALIAS.lower()}._operation = 1 THEN DELETE
+      WHEN MATCHED AND {MERGE_SOURCE_ALIAS.lower()}._operation <> 1 THEN UPDATE SET
+        {MERGE_TARGET_ALIAS.lower()}.purchase_order_id = 1
     ),
     merge_filter (
-      {MERGE_SOURCE_ALIAS}.ds > (
+      {MERGE_SOURCE_ALIAS.lower()}.ds > (
         SELECT
           MAX(ds)
         FROM db.test
       )
-      AND {MERGE_SOURCE_ALIAS}.ds > @start_ds
-      AND {MERGE_SOURCE_ALIAS}._operation <> 1
-      AND {MERGE_TARGET_ALIAS}.start_date > DATEADD(day, -7, CURRENT_DATE)
+      AND {MERGE_SOURCE_ALIAS.lower()}.ds > @start_ds
+      AND {MERGE_SOURCE_ALIAS.lower()}._operation <> 1
+      AND {MERGE_TARGET_ALIAS.lower()}.start_date > DATEADD(day, -7, CURRENT_DATE)
     ),
     batch_concurrency 1,
     forward_only FALSE,
@@ -7943,7 +8001,49 @@ def test_merge_filter():
     rendered_merge_filters = model.render_merge_filter(start="2023-01-01", end="2023-01-02")
     assert (
         rendered_merge_filters.sql()
-        == "(__MERGE_SOURCE__.ds > (SELECT MAX(ds) FROM db.test) AND __MERGE_SOURCE__.ds > '2023-01-01' AND __MERGE_SOURCE__._operation <> 1 AND __MERGE_TARGET__.start_date > DATEADD(day, -7, CURRENT_DATE))"
+        == "(__merge_source__.ds > (SELECT MAX(ds) FROM db.test) AND __merge_source__.ds > '2023-01-01' AND __merge_source__._operation <> 1 AND __merge_target__.start_date > DATEADD(day, -7, CURRENT_DATE))"
+    )
+
+
+def test_merge_filter_normalization():
+    # unquoted gets normalized
+    expressions = d.parse(
+        """
+        MODEL (
+          name db.employees,
+          kind INCREMENTAL_BY_UNIQUE_KEY (
+            unique_key name,
+            merge_filter source.salary > 0
+          )
+        );
+        SELECT 'name' AS name, 1 AS salary;
+    """
+    )
+
+    model = load_sql_based_model(expressions, dialect="snowflake")
+    assert (
+        model.merge_filter.sql(dialect="snowflake", identify=True)
+        == '"__MERGE_SOURCE__"."SALARY" > 0'
+    )
+
+    # quoted gets preserved
+    expressions = d.parse(
+        """
+        MODEL (
+          name db.employees,
+          kind INCREMENTAL_BY_UNIQUE_KEY (
+            unique_key name,
+            merge_filter source."SaLArY" > 0
+          )
+        );
+        SELECT 'name' AS name, 1 AS "SaLArY";
+    """
+    )
+
+    model = load_sql_based_model(expressions, dialect="snowflake")
+    assert (
+        model.merge_filter.sql(dialect="snowflake", identify=True)
+        == '"__MERGE_SOURCE__"."SaLArY" > 0'
     )
 
 
@@ -7970,7 +8070,7 @@ def predicate(
     )
 
     unrendered_merge_filter = (
-        f"@predicate(update_datetime) AND {MERGE_TARGET_ALIAS}.update_datetime > @start_dt"
+        f"@predicate(UPDATE_DATETIME) AND {MERGE_TARGET_ALIAS}.UPDATE_DATETIME > @start_dt"
     )
     expected_merge_filter = f"{MERGE_SOURCE_ALIAS}.UPDATE_DATETIME > DATE_ADD({MERGE_TARGET_ALIAS}.UPDATE_DATETIME, -7, 'DAY') AND {MERGE_TARGET_ALIAS}.UPDATE_DATETIME > CAST('2023-01-01 15:00:00+00:00' AS TIMESTAMPTZ)"