Rails 5 connection swapping for sharding env

multi_shard_test.rb

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