Feat: improve pretty-printing of MERGE statement (#5102)

Author: georgesittas

sqlglot/generator.py

@@ -3841,7 +3841,7 @@ def when_sql(self, expression: exp.When) -> str:
             if isinstance(then_expression.args.get("expressions"), exp.Star):
                 then = f"UPDATE {self.sql(then_expression, 'expressions')}"
             else:
-                then = f"UPDATE SET {self.expressions(then_expression, flat=True)}"
+                then = f"UPDATE SET{self.sep()}{self.expressions(then_expression)}"
         else:
             then = self.sql(then_expression)
         return f"WHEN {matched}{source}{condition} THEN {then}"

tests/dialects/test_tsql.py

@@ -962,7 +962,8 @@ def test_ddl(self):
       ProductID
   ) AS src(ProductID, OrderQty)
   ON pi.ProductID = src.ProductID
-  WHEN MATCHED AND pi.Quantity - src.OrderQty >= 0 THEN UPDATE SET pi.Quantity = pi.Quantity - src.OrderQty
+  WHEN MATCHED AND pi.Quantity - src.OrderQty >= 0 THEN UPDATE SET
+    pi.Quantity = pi.Quantity - src.OrderQty
   WHEN MATCHED AND pi.Quantity - src.OrderQty <= 0 THEN DELETE
   OUTPUT $action, Inserted.ProductID, Inserted.LocationID, Inserted.Quantity AS NewQty, Deleted.Quantity AS PreviousQty
 ) AS Changes(Action, ProductID, LocationID, NewQty, PreviousQty)

tests/fixtures/pretty.sql

@@ -449,3 +449,11 @@ SELECT
 FROM foo
 WHERE
   1 = 1 AND /* first comment */ foo.a /* second comment */ = 1;
+
+MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN UPDATE SET status = s.status, amount = s.amount;
+MERGE INTO t
+USING s
+ON t.id = s.id
+WHEN MATCHED THEN UPDATE SET
+  status = s.status,
+  amount = s.amount;