From 5938c92cc767ec825ec4ab05456857998296add4 Mon Sep 17 00:00:00 2001
From: Jeremy Evans <code@jeremyevans.net>
Date: Mon, 2 Sep 2019 15:04:52 -0700
Subject: [PATCH] Make sharding plugin integrate with server_block extension

This makes it so if you retrieve a model instance inside a
with_server block, and later save it outside the with_server block,
the model instance will remember where it was originally retrieved
from and try to save it back from the same server.

Instead of overriding Dataset#server to set the row_proc (which
won't be called when using with_server), override row_proc itself,
and have it check with the connection pool to get the current
server, which works both when using the server_block extension
and when not using it.
---
 CHANGELOG                        |  4 ++++
 lib/sequel/plugins/sharding.rb   | 16 +++++++++++-----
 spec/extensions/sharding_spec.rb |  8 ++++++++
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG b/CHANGELOG
index 802c15372c..46496e7ca7 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,3 +1,7 @@
+=== master
+
+* Make sharding plugin integrate with server_block extension (jeremyevans)
+
 === 5.24.0 (2019-09-01)
 
 * Add Database#skip_logging? private method designed for extensions to force query timing even if no logger is present (adam12) (#1640)
diff --git a/lib/sequel/plugins/sharding.rb b/lib/sequel/plugins/sharding.rb
index b0a2fe5de7..b328cec88e 100644
--- a/lib/sequel/plugins/sharding.rb
+++ b/lib/sequel/plugins/sharding.rb
@@ -107,12 +107,18 @@ module DatasetMethods
         # previous row_proc, but calls set_server on the output of that row_proc,
         # ensuring that objects retrieved by a specific shard know which shard they
         # are tied to.
-        def server(s)
-          ds = super
-          if rp = row_proc
-            ds = ds.with_row_proc(proc{|r| rp.call(r).set_server(s)})
+        def row_proc
+          rp = super
+          if rp
+            case server = db.pool.send(:pick_server, opts[:server])
+            when nil, :default, :read_only
+              # nothing
+            else
+              old_rp = rp
+              rp = proc{|r| old_rp.call(r).set_server(server)}
+            end
           end
-          ds
+          rp 
         end
       end
     end
diff --git a/spec/extensions/sharding_spec.rb b/spec/extensions/sharding_spec.rb
index 0b2987909e..92f9c9ce03 100644
--- a/spec/extensions/sharding_spec.rb
+++ b/spec/extensions/sharding_spec.rb
@@ -186,4 +186,12 @@
     ["UPDATE albums SET artist_id = 2, name = 'RF' WHERE (id = 1) -- s1", "UPDATE albums SET name = 'RF', artist_id = 2 WHERE (id = 1) -- s1"].must_include(sqls.slice!(2))
     sqls.must_equal ["SELECT * FROM albums LIMIT 1 -- s2", "BEGIN -- s1", "COMMIT -- s1"]
   end 
+
+  it "should have objects retrieved from a specific shard using with_server from server_block extension" do
+    album = @db.extension(:server_block).with_server(:s1) do
+      @Album.first
+    end
+    album.update(:name=>'MO')
+    @db.sqls.must_equal ["SELECT * FROM albums LIMIT 1 -- s1", "UPDATE albums SET name = 'MO' WHERE (id = 1) -- s1"]
+  end 
 end