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

@@ -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)

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

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