Repository: tobymao/sqlglot
Pull request: Fix!: exp.Merge condition for Trino/Postgres
Fixes #4595

Further building upon https://github.com/tobymao/sqlglot/pull/3940 there was a gap whereby the UPDATE part of a MERGE query needed to be fully qualified if it's used within a function.
The specific example we had was when we were concatenating arrays but I've substituted for a coalesce in the unit test.

Problem statement:
Trino - when matched ambiguous column when using a function
- Read dialect Trino
- Write dialect Trino
- Still exists in main

**Fully reproducible code snippet**
```sql
MERGE INTO table_a AS target USING(
  SELECT
    pk,
    my_array
  FROM table_b
) AS source ON source.pk = target.pk
WHEN MATCHED THEN UPDATE SET target.my_array = ARRAY_DISTINCT(CONCAT(source.my_array, target.my_array))
WHEN NOT MATCHED THEN INSERT (pk, my_array)
VALUES (
  source.pk,
  source.my_array
)
```

As the WHEN MATCHED is a function then this should be possible but there's a transformation applied that strips out the table aliases in the THEN part.
This then leads us to an ambiguous column error as target.my_array is transformed to just my_array.

Might a solution be to also consider if the column's parent is a function? In which case if the merge target is referenced then it is to remove ambiguity.
I'll raise a PR with such a solution shortly and will await comments.

**Official Documentation**
Please include links to official SQL documentation related to your issue.

Hints from issue discussion:

Commit message: address PR review comment

Patch to review:
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index 5afa059608..fe93f62aea 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -1570,19 +1570,25 @@ def normalize(identifier: t.Optional[exp.Identifier]) -> t.Optional[str]:
         targets.add(normalize(alias.this))
 
     for when in expression.args["whens"].expressions:
-        # only remove the target names from the THEN clause
-        # theyre still valid in the <condition> part of WHEN MATCHED / WHEN NOT MATCHED
-        # ref: https://github.com/TobikoData/sqlmesh/issues/2934
-        then = when.args.get("then")
+        # only remove the target table names from certain parts of WHEN MATCHED / WHEN NOT MATCHED
+        # they are still valid in the <condition>, the right hand side of each UPDATE and the VALUES part
+        # (not the column list) of the INSERT
+        then: exp.Insert | exp.Update | None = when.args.get("then")
         if then:
-            then.transform(
-                lambda node: (
-                    exp.column(node.this)
-                    if isinstance(node, exp.Column) and normalize(node.args.get("table")) in targets
-                    else node
-                ),
-                copy=False,
-            )
+            if isinstance(then, exp.Update):
+                for equals in then.find_all(exp.EQ):
+                    equal_lhs = equals.this
+                    if (
+                        isinstance(equal_lhs, exp.Column)
+                        and normalize(equal_lhs.args.get("table")) in targets
+                    ):
+                        equal_lhs.replace(exp.column(equal_lhs.this))
+            if isinstance(then, exp.Insert):
+                column_list = then.args.get("this")
+                if column_list is not None:
+                    for column in column_list.find_all(exp.Column):
+                        if normalize(column.args.get("table")) in targets:
+                            column.replace(exp.column(column.this))
 
     return self.merge_sql(expression)
 
diff --git a/tests/dialects/test_dialect.py b/tests/dialects/test_dialect.py
index 547bf539e7..1573329540 100644
--- a/tests/dialects/test_dialect.py
+++ b/tests/dialects/test_dialect.py
@@ -2336,6 +2336,17 @@ def test_merge(self):
             },
         )
 
+        # needs to preserve the target alias in then WHEN condition and function but not in the THEN clause
+        self.validate_all(
+            """MERGE INTO foo AS target USING (SELECT a, b FROM tbl) AS src ON src.a = target.a
+            WHEN MATCHED THEN UPDATE SET target.b = COALESCE(src.b, target.b)
+            WHEN NOT MATCHED THEN INSERT (target.a, target.b) VALUES (src.a, src.b)""",
+            write={
+                "trino": """MERGE INTO foo AS target USING (SELECT a, b FROM tbl) AS src ON src.a = target.a WHEN MATCHED THEN UPDATE SET b = COALESCE(src.b, target.b) WHEN NOT MATCHED THEN INSERT (a, b) VALUES (src.a, src.b)""",
+                "postgres": """MERGE INTO foo AS target USING (SELECT a, b FROM tbl) AS src ON src.a = target.a WHEN MATCHED THEN UPDATE SET b = COALESCE(src.b, target.b) WHEN NOT MATCHED THEN INSERT (a, b) VALUES (src.a, src.b)""",
+            },
+        )
+
     def test_substring(self):
         self.validate_all(
             "SUBSTR('123456', 2, 3)",


Code review comments:
diff_hunk: @@ -1570,19 +1570,25 @@ def normalize(identifier: t.Optional[exp.Identifier]) -> t.Optional[str]:
         targets.add(normalize(alias.this))
 
     for when in expression.args["whens"].expressions:
-        # only remove the target names from the THEN clause
-        # theyre still valid in the <condition> part of WHEN MATCHED / WHEN NOT MATCHED
-        # ref: https://github.com/TobikoData/sqlmesh/issues/2934
-        then = when.args.get("then")
+        # only remove the target table names from certain parts of WHEN MATCHED / WHEN NOT MATCHED
+        # they are still valid in the <condition>, the right hand side of each UPDATE and the VALUES part
+        # (not the column list) of the INSERT
+        then: exp.Insert | exp.Update | None = when.args.get("then")
         if then:
-            then.transform(
-                lambda node: (
-                    exp.column(node.this)
-                    if isinstance(node, exp.Column) and normalize(node.args.get("table")) in targets
-                    else node
-                ),
-                copy=False,
-            )
+            if isinstance(then, exp.Update):
+                for equals in then.find_all(exp.EQ):
+                    equal_lhs = equals.this
+                    if (
+                        isinstance(equal_lhs, exp.Column)
+                        and normalize(equal_lhs.args.get("table")) in targets
+                    ):
+                        equal_lhs.replace(exp.column(equal_lhs.this))
+            if isinstance(then, exp.Insert):
+                column_list = then.args.get("this")
+                if column_list is not None:
+                    for column in column_list.find_all(exp.Column):
path: sqlglot/dialects/dialect.py
text: @user1:
```suggestion
                column_list = then.this
                if isinstance(column_list, exp.Tuple):
                    for column in column_list.expressions:
```