Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Rails 5 connection swapping for sharding env
require "active_record"
require "minitest/autorun"
ENV["RAILS_ENV"] = "development"
ActiveRecord::Base.configurations = {
"development" => { "adapter" => "sqlite3", "database" => "fake.db" },
"shard1" => { "adapter" => "sqlite3", "database" => "shard1.db" },
"shard2" => { "adapter" => "sqlite3", "database" => "shard2.db" }
}
ActiveRecord::Base.establish_connection
# this needs to be added to a initializer
ActiveRecord::Base.connection_handler.establish_connection :shard1
ActiveRecord::Base.connection_handler.establish_connection :shard2
module Shard
extend self
mattr_accessor :current_shard
def in_shard(n)
prev_shard, self.current_shard = self.current_shard, n
yield
ensure
self.current_shard = prev_shard
end
def in_all_shards(&blk)
[1, 2].each do |n|
in_shard(n, &blk)
end
end
end
class ApplicationRecord < ActiveRecord::Base
def self.connection_specification_name
if Shard.current_shard.nil?
raise "woops"
end
"shard#{Shard.current_shard}"
end
end
class Foo < ApplicationRecord
Shard.in_all_shards do
connection.create_table table_name, force: true do |t|
t.string :name
end
end
end
class ConnHandlerSwapTest < Minitest::Test
def test_insert_in_shards
Shard.in_shard(1) do
Foo.create!(name: 'arthurnn')
assert_equal 1, Foo.where(name: 'arthurnn').count
end
Shard.in_shard(2) do
assert_equal 0, Foo.where(name: 'arthurnn').count
Foo.create!(name: 'arthurnn2')
assert_equal 1, Foo.where(name: 'arthurnn2').count
end
assert_raises do
# not shard selected
Foo.where(name: 'arthurnn2').count
end
end
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.