Repository: tobymao/sqlglot

Issue:
Support Postgres 17 MERGE Returning
**Is your feature request related to a problem? Please describe.**
Postgres 17, which is likely going to be generally available later this month adds the ability for MERGE statements to return rows. However, sqlglot is unable to parse these.

**Describe the solution you'd like**
Take this example, [from this blog](https://neon.tech/blog/top-3-features-in-postgres-17):

```sql
MERGE INTO hero h
USING (VALUES ('Wade', 'Wilson', 'Deadpool')) v(first_name, last_name, hero_name)
ON h.hero_name = v.hero_name
WHEN MATCHED THEN
  UPDATE SET first_name = v.first_name, last_name = v.last_name
WHEN NOT MATCHED THEN
  INSERT (first_name, last_name, hero_name)
  VALUES (v.first_name, v.last_name, v.hero_name)
RETURNING merge_action(), *;
```

Trying to parse this yields:

```
Invalid expression / Unexpected token. Line 10, Col: 9.
   THEN
  INSERT (first_name, last_name, hero_name)
  VALUES (v.first_name, v.last_name, v.hero_name)
RETURNING merge_action(), *;
```



Discussion:

Patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 6bbd0fb4..21d3a8fd 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -6279,6 +6279,7 @@ class Merge(Expression):
         "on": True,
         "expressions": True,
         "with": False,
+        "returning": False,
     }
 
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 9a346b24..6e11d68c 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -3587,9 +3587,12 @@ class Generator(metaclass=_Generator):
         on = f"ON {self.sql(expression, 'on')}"
         expressions = self.expressions(expression, sep=" ", indent=False)
         sep = self.sep()
+        returning = self.expressions(expression, key="returning", indent=False)
+        returning = f"RETURNING {returning}" if returning else ""
 
         return self.prepend_ctes(
-            expression, f"MERGE INTO {this}{table_alias}{sep}{using}{sep}{on}{sep}{expressions}"
+            expression,
+            f"MERGE INTO {this}{table_alias}{sep}{using}{sep}{on}{sep}{expressions}{sep}{returning}",
         )
 
     def tochar_sql(self, expression: exp.ToChar) -> str:
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index a15dd815..840f2f13 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -6825,6 +6825,7 @@ class Parser(metaclass=_Parser):
             using=using,
             on=on,
             expressions=self._parse_when_matched(),
+            returning=self._match(TokenType.RETURNING) and self._parse_csv(self._parse_bitwise),
         )
 
     def _parse_when_matched(self) -> t.List[exp.When]:
@@ -6865,7 +6866,7 @@ class Parser(metaclass=_Parser):
             elif self._match(TokenType.DELETE):
                 then = self.expression(exp.Var, this=self._prev.text)
             else:
-                then = None
+                then = self._parse_var_from_options(self.CONFLICT_ACTIONS)
 
             whens.append(
                 self.expression(


Test patch:
diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
index f3f21a9b..1b89ecd1 100644
--- a/tests/dialects/test_postgres.py
+++ b/tests/dialects/test_postgres.py
@@ -794,6 +794,9 @@ class TestPostgres(Validator):
         self.validate_identity(
             'SELECT js, js IS JSON ARRAY WITH UNIQUE KEYS AS "array w. UK?", js IS JSON ARRAY WITHOUT UNIQUE KEYS AS "array w/o UK?", js IS JSON ARRAY UNIQUE KEYS AS "array w UK 2?" FROM t'
         )
+        self.validate_identity(
+            "MERGE INTO target_table USING source_table AS source ON target.id = source.id WHEN MATCHED THEN DO NOTHING WHEN NOT MATCHED THEN DO NOTHING RETURNING MERGE_ACTION(), *"
+        )
 
     def test_ddl(self):
         # Checks that user-defined types are parsed into DataType instead of Identifier