From 326df130bb5f5fd6d85ab966e534794bcb5b0a4a Mon Sep 17 00:00:00 2001
From: Jeremy Evans <code@jeremyevans.net>
Date: Fri, 7 Jun 2024 12:09:08 -0700
Subject: [PATCH] Support MERGE RETURNING on PostgreSQL 17+

---
 CHANGELOG                              |   2 +
 lib/sequel/adapters/shared/postgres.rb |  25 +++++-
 spec/adapters/postgres_spec.rb         | 102 +++++++++++++++++++++++++
 3 files changed, 127 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG b/CHANGELOG
index c37b35d6a..edda30d2f 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,5 +1,7 @@
 === master
 
+* Support MERGE RETURNING on PostgreSQL 17+ (jeremyevans)
+
 * Remove use of logger library in bin/sequel (jeremyevans)
 
 * Support :connect_opts_proc Database option for late binding options (jeremyevans) (#2164)
diff --git a/lib/sequel/adapters/shared/postgres.rb b/lib/sequel/adapters/shared/postgres.rb
index ee118c5cc..c147d7de8 100644
--- a/lib/sequel/adapters/shared/postgres.rb
+++ b/lib/sequel/adapters/shared/postgres.rb
@@ -2058,6 +2058,16 @@ def lock(mode, opts=OPTS)
         nil
       end
 
+      # Support MERGE RETURNING on PostgreSQL 17+.
+      def merge(&block)
+        sql = merge_sql
+        if uses_returning?(:merge)
+          returning_fetch_rows(sql, &block)
+        else
+          execute_ddl(sql)
+        end
+      end
+
       # Return a dataset with a WHEN MATCHED THEN DO NOTHING clause added to the
       # MERGE statement.  If a block is passed, treat it as a virtual row and
       # use it as additional conditions for the match.
@@ -2170,9 +2180,14 @@ def supports_nowait?
         true
       end
 
-      # Returning is always supported.
+      # MERGE RETURNING is supported on PostgreSQL 17+. Other RETURNING is supported
+      # on all supported PostgreSQL versions.
       def supports_returning?(type)
-        true
+        if type == :merge
+          server_version >= 170000
+        else
+          true
+        end
       end
 
       # PostgreSQL supports pattern matching via regular expressions
@@ -2295,6 +2310,12 @@ def _merge_matched_sql(sql, data)
       end
       alias _merge_not_matched_sql _merge_matched_sql
 
+      # Support MERGE RETURNING on PostgreSQL 17+.
+      def _merge_when_sql(sql)
+        super
+        insert_returning_sql(sql) if uses_returning?(:merge)
+      end
+
       # Format TRUNCATE statement with PostgreSQL specific options.
       def _truncate_sql(table)
         to = @opts[:truncate_opts] || OPTS
diff --git a/spec/adapters/postgres_spec.rb b/spec/adapters/postgres_spec.rb
index 3cc8cde86..9307f30f0 100644
--- a/spec/adapters/postgres_spec.rb
+++ b/spec/adapters/postgres_spec.rb
@@ -5928,3 +5928,105 @@ def c2.name; 'Bar' end
     end
   end
 end
+
+
+describe "MERGE RETURNING" do
+  before(:all) do
+    @db = DB
+    @db.create_table!(:m1){Integer :i1; Integer :a}
+    @db.create_table!(:m2){Integer :i2; Integer :b}
+    @m1 = @db[:m1].returning
+    @m2 = @db[:m2]
+  end
+  after do
+    @m1.delete
+    @m2.delete
+  end
+  after(:all) do
+    @db.drop_table?(:m1, :m2)
+  end
+
+  def merge(ds)
+    a = []
+    ds.merge{|h| a << h}
+    a
+  end
+
+  def check(ds)
+    @m2.insert(1, 2)
+
+    @m1.all.must_equal []
+
+    # INSERT
+    merge(ds).must_equal [{:i2=>1, :b=>2, :i1=>1, :a=>13}]
+    @m1.all.must_equal [{:i1=>1, :a=>13}]
+
+    # UPDATE
+    merge(ds).must_equal [{:i2=>1, :b=>2, :i1=>12, :a=>35}]
+    @m1.all.must_equal [{:i1=>12, :a=>35}]
+
+    # UPDATE with specific RETURNING columns
+    @m1.update(:i1=>1, :a=>13)
+    merge(ds.returning(:b)).must_equal [{:b=>2}]
+    @m1.all.must_equal [{:i1=>12, :a=>35}]
+
+    # DELETE MATCHING current row, INSERT NOT MATCHED new row
+    @m2.insert(12, 3)
+    merge(ds).must_equal [{:i2=>1, :b=>2, :i1=>1, :a=>13}, {:i2=>12, :b=>3, :i1=>12, :a=>35}]
+    @m1.all.must_equal [{:i1=>1, :a=>13}]
+
+    # MATCHED DO NOTHING
+    @m2.where(:i2=>12).delete
+    @m1.update(:a=>51)
+    merge(ds).must_equal []
+    @m1.all.must_equal [{:i1=>1, :a=>51}]
+
+    # NOT MATCHED DO NOTHING
+    @m1.delete
+    @m2.update(:b=>51)
+    merge(ds).must_equal []
+    @m1.all.must_equal []
+  end
+  
+  it "should allow inserts, updates, and deletes based on conditions in a single MERGE statement" do
+    ds = @m1.
+      merge_using(:m2, :i1=>:i2).
+      merge_insert(:i1=>Sequel[:i2], :a=>Sequel[:b]+11){b <= 50}.
+      merge_delete{{:a => 30..50}}.
+      merge_update(:i1=>Sequel[:i1]+:i2+10, :a=>Sequel[:a]+:b+20){a <= 50}
+    check(ds)
+  end
+
+  it "should support WITH clauses" do
+    ds = @m1.
+      with(:m3, @db[:m2]).
+      merge_using(:m3, :i1=>:i2).
+      merge_insert(:i1=>Sequel[:i2], :a=>Sequel[:b]+11){b <= 50}.
+      merge_delete{{:a => 30..50}}.
+      merge_update(:i1=>Sequel[:i1]+:i2+10, :a=>Sequel[:a]+:b+20){a <= 50}
+    check(ds)
+  end if DB.dataset.supports_cte?
+
+  it "should support inserts with just columns" do
+    ds = @m1.
+      merge_using(:m2, :i1=>:i2).
+      merge_insert(Sequel[:i2], Sequel[:b]+11){b <= 50}.
+      merge_delete{{:a => 30..50}}.
+      merge_update(:i1=>Sequel[:i1]+:i2+10, :a=>Sequel[:a]+:b+20){a <= 50}
+    check(ds)
+  end
+
+  it "should calls inserts, updates, and deletes without conditions" do
+    @m2.insert(1, 2)
+    ds = @m1.merge_using(:m2, :i1=>:i2)
+    
+    merge(ds.merge_insert(:i2, :b)).must_equal [{:i2=>1, :b=>2, :i1=>1, :a=>2}]
+    @m1.all.must_equal [{:i1=>1, :a=>2}]
+
+    merge(ds.merge_update(:a=>Sequel[:a]+1)).must_equal [{:i2=>1, :b=>2, :i1=>1, :a=>3}]
+    @m1.all.must_equal [{:i1=>1, :a=>3}]
+
+    merge(ds.merge_delete).must_equal [{:i2=>1, :b=>2, :i1=>1, :a=>3}]
+    @m1.all.must_equal []
+  end
+end if DB.server_version >= 170000