rails-sharding 0.1.1 → 1.0.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 05fd12669243e14c2d6f1b1748acf461b8316dbb
-  data.tar.gz: d1508fddee17b0a5c7eb0203b8dd1c9b52b74c66
+  metadata.gz: 58f0c90d4f03a5f3aa3d4ff363cdac165c948f8f
+  data.tar.gz: d927face94cf2c001534c1353b322294ab1984b8
 SHA512:
-  metadata.gz: 5ea75e62463c090f6806baedfd648768007c192d4e3a9e162c162083eb015cacb8eb6dfbc8ed0ce739bb9cc17c9a3889abc6b0bd02bb24f3f76470478996e85c
-  data.tar.gz: 524dfe1e8c7ea123209746accfdb726ea0a6011f43c460cbcf2cff3c11531b2ee7afa99c26091d54ad8792dd05dca62aff2badaee5290babb1f72631ed0e94d0
+  metadata.gz: aa192b584a7ef8e58be570798c8d35ae04f17134791ed325a272f32dc5c32d3d847e60d06894da35093736d7c29f52bba32414b4d27b6d5107e2390e86c43221
+  data.tar.gz: cd329c1e3473f0f8a11f93bfc15bdd5278ebb041b8fccf6d5e2b7b496fa229919914b4080e4d97913e61ad090c6b67dfdf2a9c22e729cace4129e4f63e92bf95

data/README.md

@@ -4,8 +4,9 @@
 [![Code Climate](https://codeclimate.com/github/hsgubert/rails-sharding/badges/gpa.svg)](https://codeclimate.com/github/hsgubert/rails-sharding)
 [![Test Coverage](https://codeclimate.com/github/hsgubert/rails-sharding/badges/coverage.svg)](https://codeclimate.com/github/hsgubert/rails-sharding/coverage)
 [![Gem Version](https://badge.fury.io/rb/rails-sharding.svg)](https://badge.fury.io/rb/rails-sharding)
+[![Dependency Status](https://gemnasium.com/badges/github.com/hsgubert/rails-sharding.svg)](https://gemnasium.com/github.com/hsgubert/rails-sharding)
 
-Simple and robust sharding for Rails, including Migrations and ActiveRecord extensions
+Simple and robust sharding gem for Rails, including Migrations and ActiveRecord extensions
 
 This gems allows you to easily create extra databases to your rails application, and freely allocate ActiveRecord instances to any of the databases. It also provides rake tasks and migrations to help you manage the schema by shard groups.
 
@@ -27,7 +28,10 @@ You can also use the block syntax, where all your queries inside will be directe
 You can also pick and choose which models will be shardable, so that all the models that are not shardable will still be retrieved from the master database, even if inside a using_shard block.
 
 ## Compatibility
-As of now this gem has been tested only with Rails 4.2. It does not work yet with Rails 5.
+Gem version 1.x.x -> compatible with Rails 5.0
+
+Gem version 0.1.1 -> compatible with Rails 4.2
+
 
 ## Installation
 
@@ -83,7 +87,7 @@ rake shards:create
 Go to the directory `db/shards_migrations/shard_group1` and add all migrations that you want to run on the shards of `shard_group1`. By design, all shards in a same group should always have the same schema. For example, add the following migration to your `db/shards_migrations/shard_group1`:
 ```ruby
 # 20160808000000_create_users.rb
-class CreateClients < ActiveRecord::Migration
+class CreateClients < ActiveRecord::Migration[5.0]
   def up
     create_table :users do |t|
       t.string :username, :limit => 100
@@ -140,6 +144,29 @@ rake shards:migrate SHARD_GROUP=shard_group_1
 rake shards:migrate SHARD_GROUP=shard_group_1 SHARD=shard1
 ```
 
+## Gem Options
+Running the `rails g rails_sharding:scaffold` will create an initializer at `config/initializers/rails-sharding.rb`. You can pass additional configurations on this initializer to control the gem behavior. You can see below all available options and their default values:
+```ruby
+# config/initializers/rails-sharding.rb
+
+Rails::Sharding.setup do |config|
+  # If true one connection will be established per shard (in every shard group) on startup.
+  # If false the user must call Shards::ConnectionHandler.establish_connection(shard_group, shard_name) manually at least once before using each shard.
+  config.establish_all_connections_on_setup = true
+
+  # If true the method #using_shard will be mixed in ActiveRecord scopes. Put this to false if you don't want the gem to modify ActiveRecord
+  config.extend_active_record_scope = true,
+
+  # Specifies where to find the definition of the shards configurations
+  config.shards_config_file = 'config/shards.yml',
+
+  # Specifies where to find the migrations for each shard group
+  config.shards_migrations_dir = 'db/shards_migrations',
+
+  # Specifies where to find the schemas for each shard group
+  config.shards_schemas_dir = 'db/shards_schemas'
+end
+```
 
 ## Development and Contributing
 
@@ -154,4 +181,4 @@ The gem is available as open source under the terms of the [MIT License](http://
 
 ## Acknowledgements
 
-This gem was inspired and based on several other gems like: [octopus](https://github.com/thiagopradi/octopus), [shard_handler](https://github.com/locaweb/shard_handler) and [active_record_shards](https://github.com/zendesk/active_record_shards).
+This gem was inspired on several other gems like: [octopus](https://github.com/thiagopradi/octopus), [shard_handler](https://github.com/locaweb/shard_handler) and [active_record_shards](https://github.com/zendesk/active_record_shards).

data/Rakefile

@@ -6,7 +6,8 @@ RSpec::Core::RakeTask.new(:spec)
 task :default => :spec
 
 # defines an environment task so we can run rake tasks from lib/tasks/rails-sharding.rake.
-# They require the :environment task, which is defined by rails, but we don need anything
+# The tasks on rails-sharding.rake depend on the :environment task, which is usuallu defined
+# by rails. In our case, we just stub it so the rake tasks run
 task :environment do
   # do nothing
 end
@@ -16,12 +17,7 @@ namespace :db do
 
     desc 'Loads gem test environment and rake tasks from gem'
     task :load_env do
-      require 'rspec'
-
-      # requires spec helper but ensures no test coverage is reported to codeclimate
-      ENV['CODECLIMATE_REPO_TOKEN'] = nil
-      require './spec/spec_helper'
-      
+      require './spec/load_gem_test_env'
       load 'lib/tasks/rails-sharding.rake'
     end
 

data/lib/generators/templates/shards.yml.example

@@ -1,7 +1,7 @@
 default: &default
   adapter: mysql2
   encoding: utf8
-  reconnect: false
+  reconnect: true
   pool: 5
   username: ___
   password: ___

data/lib/rails/sharding/config.rb

@@ -13,6 +13,16 @@ module Rails::Sharding
       # this to false if you don't want the gem to modify ActiveRecord
       extend_active_record_scope: true,
 
+      # If true the query logs of ActiveRecord will be tagged with the corresponding
+      # shard you're querying
+      add_shard_tag_to_query_logs: true,
+
+      # If true a warning will be printed everytime a using_shard block ends without
+      # the shard connection being retrieved at least once inside the block. This warning
+      # is helpful to remember the developer to include ShardableModel module on the
+      # sharded models, otherwise they will always connect to the master database.
+      no_connection_retrieved_warning: true,
+
       # Specifies where to find the definition of the shards configurations
       shards_config_file: 'config/shards.yml',
 
@@ -27,6 +37,5 @@ module Rails::Sharding
       self.cattr_accessor config_name
       self.send(config_name.to_s + '=', default_value)
     end
-
   end
 end

data/lib/rails/sharding/connection_handler.rb

@@ -17,48 +17,64 @@ module Rails::Sharding
     def self.establish_connection(shard_group, shard_name, environment=nil)
       self.setup unless defined? @@connection_handler
 
-      unless configurations = (environment.nil? ? Core.configurations : Core.configurations(environment))
+      configurations = (environment.nil? ? Core.configurations : Core.configurations(environment))
+      if configurations.nil?
         raise Errors::ConfigNotFoundError, "Cannot find configuration for environment '#{environment}' in #{Config.shards_config_file}"
       end
 
-      unless shard_group_configurations = configurations[shard_group.to_s]
+      shard_group_configurations = configurations[shard_group.to_s]
+      if shard_group_configurations.nil?
         raise Errors::ConfigNotFoundError, "Cannot find configuration for shard_group '#{shard_group}' in environment '#{environment}' in #{Config.shards_config_file}"
       end
 
       resolver = ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new(shard_group_configurations)
       begin
-        connection_spec = resolver.spec(shard_name.to_sym)
-      rescue ActiveRecord::AdapterNotSpecified => e
+        connection_name = connection_name(shard_group, shard_name)
+        connection_spec = resolver.spec(shard_name.to_sym, connection_name)
+      rescue ActiveRecord::AdapterNotSpecified
         raise Errors::ConfigNotFoundError, "Cannot find configuration for shard '#{shard_group}:#{shard_name}' in environment '#{environment}' in #{Config.shards_config_file}"
       end
 
-      # since rails requires a class to be the connection owner, we trick rails passing
-      # an instance of the ConnectionPoolOwner class, that responds to the #name method
-      connection_handler.establish_connection(connection_pool_owner(shard_group, shard_name), connection_spec)
+      connection_handler.establish_connection(connection_spec)
     end
 
     def self.connection_pool(shard_group, shard_name)
-      connection_handler.retrieve_connection_pool(connection_pool_owner(shard_group, shard_name))
-    rescue Errors::ConnectionPoolRetrievalError
+      if connection_pool = connection_handler.retrieve_connection_pool(connection_name(shard_group, shard_name))
+        return connection_pool
+      end
+
       # mimicking behavior of rails at:
-      # https://github.com/rails/rails/blob/4-2-stable/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb#507
-      raise ActiveRecord::ConnectionNotEstablished, "No connection pool for shard #{connection_name(shard_group, shard_name)}"
+      # https://github.com/rails/rails/blob/v5.0.0.1/activerecord/lib/active_record/connection_handling.rb#124
+      raise ActiveRecord::ConnectionNotEstablished, "No connection pool for shard #{connection_name(shard_group, shard_name)}" if connection_pool.nil?
     end
 
     def self.retrieve_connection(shard_group, shard_name)
-      connection_handler.retrieve_connection(connection_pool_owner(shard_group, shard_name))
+      connection_name = connection_name(shard_group, shard_name)
+      connection = connection_handler.retrieve_connection(connection_name)
+
+      if connection && Config.add_shard_tag_to_query_logs
+        add_shard_tag_to_connection_log(connection, connection_name)
+      else
+        connection
+      end
     end
 
     def self.connected?(shard_group, shard_name)
-      connection_handler.connected?(connection_pool_owner(shard_group, shard_name))
+      connection_handler.connected?(connection_name(shard_group, shard_name))
     end
 
     def self.with_connection(shard_group, shard_name, &block)
-    	connection_pool(shard_group, shard_name).with_connection(&block)
+    	connection_pool(shard_group, shard_name).with_connection do |connection|
+        if connection && Config.add_shard_tag_to_query_logs
+          connection_name = connection_name(shard_group, shard_name)
+          add_shard_tag_to_connection_log(connection, connection_name)
+        end
+        block.call(connection)
+      end
     end
 
     def self.remove_connection(shard_group, shard_name)
-      connection_handler.remove_connection(connection_pool_owner(shard_group, shard_name))
+      connection_handler.remove_connection(connection_name(shard_group, shard_name))
     end
 
   private
@@ -70,12 +86,6 @@ module Rails::Sharding
 
     def self.setup
       @@connection_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
-      @@connection_pool_owners = {}
-    end
-
-    def self.connection_pool_owner(shard_group, shard_name)
-      connection_name = self.connection_name(shard_group, shard_name)
-      @@connection_pool_owners[connection_name] ||= ConnectionPoolOwner.new(connection_name)
     end
 
     # Assembles connection name in the format "shard_group:shard_name"
@@ -83,22 +93,30 @@ module Rails::Sharding
       shard_group.to_s + ':' + shard_name.to_s
     end
 
-    class ConnectionPoolOwner
-      attr_reader :name
-
-      def initialize(name)
-      	@name = name
+    # Adds a shard tag to the log of all queries executed through this connection
+    def self.add_shard_tag_to_connection_log(connection, shard_tag)
+      # avoids modifing connection twice
+      if connection.respond_to? :shard_tag
+        connection.shard_tag = shard_tag
+        return connection
       end
 
-      # Safeguard in case pool cannot be retrieved for owner. This makes the error clear
-      def superclass
-      	raise Errors::ConnectionPoolRetrievalError, "ConnectionPool could not be retrieved for #{self}. See https://github.com/rails/rails/blob/4-2-stable/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb#607"
-      end
+      # creates #shard_tag attribute in connection
+      connection.singleton_class.send(:attr_accessor, :shard_tag)
+      connection.shard_tag = shard_tag
 
-      # in case owner ends up printed by rails in an error message when retrieving connection
-      def to_s
-        "ConnectionPoolOwner with name #{self.name}"
+      # create an alias #original_execute, as a copy of the #execute for this connection
+      connection.singleton_class.send(:alias_method, :original_execute, :execute)
+
+      # defines a new #execute that adds a tag to the log
+      class << connection
+        def execute(sql, name=nil)
+          name = (name.to_s + " (#{shard_tag})").strip
+          self.original_execute(sql, name)
+        end
       end
+
+      connection
     end
   end
 end

data/lib/rails/sharding/core.rb

@@ -10,23 +10,30 @@ module Rails::Sharding
 
     # Opens a block where all queries will be directed to the selected shard
     def self.using_shard(shard_group, shard_name)
-      raise 'Cannot nest using_shard blocks' if ShardThreadRegistry.connecting_to_shard?
-
-      ShardThreadRegistry.current_shard_group = shard_group
-      ShardThreadRegistry.current_shard_name = shard_name
+      ShardThreadRegistry.push_current_shard(shard_group, shard_name)
       yield
     ensure
+      shard_group, shard_name, connection_used = ShardThreadRegistry.pop_current_shard
+
+      # shows warning to user
+      if Config.no_connection_retrieved_warning && !connection_used
+        puts "Warning: no connection to shard '#{shard_group}:#{shard_name}' was retrieved inside the using_shard block. Make sure you don't forget to include Rails::Sharding::ShardableModel to the models you want to be sharded. Disable this warning with Rails::Sharding::Config.no_connection_retrieved_warning = false."
+      end
+
       # Releases connections in case user left some connection in the reserved state
       # (by calling retrieve_connection instead of with_connection). Also, using
       # normal activerecord queries leaves a connection in the reserved state
-      ConnectionHandler.connection_pool(*ShardThreadRegistry.current_shard_group_and_name).release_connection
-      ShardThreadRegistry.connect_back_to_master!
+      # Obs: don't do this with a master database connection
+      ConnectionHandler.connection_pool(shard_group, shard_name).release_connection if shard_group && shard_name
     end
 
     def self.configurations(environment=Rails.env)
       @@db_configs ||= YAML.load_file(Config.shards_config_file)
-      @@db_configs[environment]
-    rescue Errno::ENOENT => e
+      environment_config = @@db_configs[environment]
+      return environment_config if environment_config
+
+      raise Errors::ConfigNotFoundError, 'Found no shard configurations for enviroment "' + environment + '" in ' + Config.shards_config_file.to_s + ' file was not found'
+    rescue Errno::ENOENT
       raise Errors::ConfigNotFoundError, Config.shards_config_file.to_s + ' file was not found'
     end
 
@@ -46,6 +53,23 @@ module Rails::Sharding
       self.configurations[shard_group.to_s].keys
     end
 
+    # yields a block for each shard in each shard group, with its configurations
+    # shard_group_filter: if passed yields only shards of this group
+    # shard_name_filter: if passed yields only shards with this name
+    def self.for_each_shard(shard_group_filter=nil, shard_name_filter=nil)
+      shard_group_filter.to_s if shard_group_filter
+      shard_name_filter.to_s if shard_name_filter
+
+      configurations.each do |shard_group, shards_configurations|
+        next if shard_group_filter && shard_group_filter != shard_group.to_s
+
+        shards_configurations.each do |shard, configuration|
+          next if shard_name_filter && shard_name_filter != shard.to_s
+          yield shard_group, shard, configuration
+        end
+      end
+    end
+
     # Method that should be called on a rails initializer
     def self.setup
       if block_given?

data/lib/rails/sharding/shard_thread_registry.rb

@@ -1,11 +1,21 @@
-require 'active_support/per_thread_registry'
 
 module Rails::Sharding
   class ShardThreadRegistry
-    # creates two thread-specific variables
-    extend ActiveSupport::PerThreadRegistry
-    attr_accessor :_current_shard_group
-    attr_accessor :_current_shard_name
+    # Creates two thread-specific stacks to store the shard of connection
+    # The top of the stack indicates the current connection
+    # This allows us to have nested blocks of #using_shard and keep track of the
+    # connections as we open/close those blocks
+    thread_mattr_accessor :_shard_group_stack
+    thread_mattr_accessor :_shard_name_stack
+
+    # auxiliary stack that keeps track of wether each shard connection was used
+    # inside its respective using_shard block (so we can print an alert if not)
+    thread_mattr_accessor :_shard_connection_used_stack
+
+    # accessors that initialize stacks if necessary
+    def self.shard_group_stack; self._shard_group_stack ||= [] end;
+    def self.shard_name_stack; self._shard_name_stack ||= [] end;
+    def self.shard_connection_used_stack; self._shard_connection_used_stack ||= [] end;
 
     def self.connecting_to_master?
       current_shard_group.nil? || current_shard_name.nil?
@@ -15,29 +25,48 @@ module Rails::Sharding
       !connecting_to_master?
     end
 
+    # Clears the connection stack and goes back to connecting to master
     def self.connect_back_to_master!
-      self.current_shard_group = nil
-      self.current_shard_name = nil
+      shard_group_stack.clear
+      shard_name_stack.clear
+      shard_connection_used_stack.clear
     end
 
     # Returns the current shard group (for the current Thread)
     def self.current_shard_group
-      _current_shard_group
-    end
-
-    # Sets the current shard group (for the current Thread)
-    def self.current_shard_group=(group)
-      self._current_shard_group = group.blank? ? nil : group.to_sym
+      shard_group_stack.last
     end
 
     # Returns the current shard name (for the current Thread)
     def self.current_shard_name
-      _current_shard_name
+      shard_name_stack.last
+    end
+
+    def self.current_connection_used?
+      shard_connection_used_stack.last
+    end
+
+    # adds shard connection to the stack
+    def self.push_current_shard(group, name)
+      # this line supresses the unused connection warning when there are nested
+      # using_shard blocks. We suppress the warning because we view nested using_shard
+      # blocks as a override
+      notify_connection_retrieved
+
+      shard_group_stack.push(group.blank? ? nil : group.to_sym)
+      shard_name_stack.push(name.blank? ? nil : name.to_sym)
+      shard_connection_used_stack.push(false)
+    end
+
+    # notifies the current connection was used (wee keep track of this to warn
+    # the user in case the connection is not used)
+    def self.notify_connection_retrieved
+      shard_connection_used_stack[-1] = true if shard_connection_used_stack.present?
     end
 
-    # Sets the current shard name (for the current Thread)
-    def self.current_shard_name=(name)
-      self._current_shard_name = name.blank? ? nil : name.to_sym
+    # removes shard connection to the stack
+    def self.pop_current_shard
+      [shard_group_stack.pop, shard_name_stack.pop, shard_connection_used_stack.pop]
     end
 
     def self.current_shard_group_and_name

data/lib/rails/sharding/shardable_model.rb

@@ -32,6 +32,8 @@ module Rails::Sharding
 
       # @overrides ActiveRecord::ConnectionHandling#connection_pool
       def sharded_connection_pool
+        ShardThreadRegistry.notify_connection_retrieved
+
         if ShardThreadRegistry.connecting_to_master?
           return original_connection_pool
         else
@@ -41,6 +43,8 @@ module Rails::Sharding
 
       # @overrides ActiveRecord::ConnectionHandling#retrieve_connection
       def sharded_retrieve_connection
+        ShardThreadRegistry.notify_connection_retrieved
+
         if ShardThreadRegistry.connecting_to_master?
           return original_retrieve_connection
         else

data/lib/rails/sharding/version.rb

@@ -1,5 +1,5 @@
 module Rails
   module Sharding
-    VERSION = "0.1.1"
+    VERSION = "1.0.1"
   end
 end

data/lib/tasks/rails-sharding.rake

@@ -8,29 +8,36 @@ shards_namespace = namespace :shards do
     ActiveRecord::Base.include(Rails::Sharding::ShardableModel) unless ActiveRecord::Base.ancestors.include? Rails::Sharding::ShardableModel
   end
 
-  desc "Creates database shards (options: RAILS_ENV=x SHARD_GROUP=x SHARD=x)"
-  task create: [:environment] do
+  # for each of the shards, check that 1) the environment set in the ar_internal_metadata
+  # table matches the current rails env and 2) it is not a protected environment
+  # (defined in ActiveRecord::Base.protected_environments)
+  desc "Checks if the environment is not protected and if the shards match the current environment (options: RAILS_ENV=x SHARD_GROUP=x SHARD=x)"
+  task check_protected_environments: [:_make_activerecord_base_shardable] do
     Rails::Sharding.configurations.each do |shard_group, shards_configurations|
       next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
 
-      shards_configurations.each do |shard, configuration|
+      shards_configurations.each do |shard, _|
         next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
-        puts "== Creating shard #{shard_group}:#{shard}"
-        ActiveRecord::Tasks::DatabaseTasks.create(configuration)
+        Rails::Sharding.using_shard(shard_group, shard) do
+          ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+        end
       end
     end
   end
 
-  desc "Drops database shards (options: RAILS_ENV=x SHARD_GROUP=x SHARD=x)"
-  task drop: [:environment] do
-    Rails::Sharding.configurations.each do |shard_group, shards_configurations|
-      next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
+  desc "Creates database shards (options: RAILS_ENV=x SHARD_GROUP=x SHARD=x)"
+  task create: [:environment] do
+    Rails::Sharding.for_each_shard(ENV["SHARD_GROUP"], ENV["SHARD"]) do |shard_group, shard, configuration|
+      puts "== Creating shard #{shard_group}:#{shard}"
+      ActiveRecord::Tasks::DatabaseTasks.create(configuration)
+    end
+  end
 
-      shards_configurations.each do |shard, configuration|
-        next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
-        puts "== Dropping shard #{shard_group}:#{shard}"
-        ActiveRecord::Tasks::DatabaseTasks.drop(configuration)
-      end
+  desc "Drops database shards (options: RAILS_ENV=x SHARD_GROUP=x SHARD=x)"
+  task drop: [:environment, :check_protected_environments] do
+    Rails::Sharding.for_each_shard(ENV["SHARD_GROUP"], ENV["SHARD"]) do |shard_group, shard, configuration|
+      puts "== Dropping shard #{shard_group}:#{shard}"
+      ActiveRecord::Tasks::DatabaseTasks.drop(configuration)
     end
   end
 
@@ -40,11 +47,9 @@ shards_namespace = namespace :shards do
       next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
 
       # configures path for migrations of this shard group and creates dir if necessary
-      shard_group_migrations_dir = File.join(Rails::Sharding::Config.shards_migrations_dir, shard_group.to_s)
-      ActiveRecord::Tasks::DatabaseTasks.migrations_paths = shard_group_migrations_dir
-      FileUtils.mkdir_p(shard_group_migrations_dir)
+      setup_migrations_path(shard_group)
 
-      shards_configurations.each do |shard, configuration|
+      shards_configurations.each do |shard, _|
         next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
         puts "== Migrating shard #{shard_group}:#{shard}"
         Rails::Sharding.using_shard(shard_group, shard) do
@@ -78,18 +83,13 @@ shards_namespace = namespace :shards do
     task dump: [:_make_activerecord_base_shardable] do
       require "active_record/schema_dumper"
 
-      Rails::Sharding.configurations.each do |shard_group, shards_configurations|
-        next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
+      Rails::Sharding.for_each_shard(ENV["SHARD_GROUP"], ENV["SHARD"]) do |shard_group, shard, _configuration|
+        puts "== Dumping schema of #{shard_group}:#{shard}"
 
-        shards_configurations.each do |shard, configuration|
-          next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
-          puts "== Dumping schema of #{shard_group}:#{shard}"
-
-          schema_filename = shard_schema_path(shard_group, shard)
-          File.open(schema_filename, "w:utf-8") do |file|
-            Rails::Sharding.using_shard(shard_group, shard) do
-              ActiveRecord::SchemaDumper.dump(Rails::Sharding::ConnectionHandler.retrieve_connection(shard_group, shard), file)
-            end
+        schema_filename = shard_schema_path(shard_group, shard)
+        File.open(schema_filename, "w:utf-8") do |file|
+          Rails::Sharding.using_shard(shard_group, shard) do
+            ActiveRecord::SchemaDumper.dump(Rails::Sharding::ConnectionHandler.retrieve_connection(shard_group, shard), file)
           end
         end
       end
@@ -100,13 +100,14 @@ shards_namespace = namespace :shards do
     end
 
     desc "Loads schema.rb file into the shards (options: RAILS_ENV=x, SHARD_GROUP=x, SHARD=x)"
-    task load: [:_make_activerecord_base_shardable] do
+    task load: [:_make_activerecord_base_shardable, :check_protected_environments] do
       Rails::Sharding.configurations.each do |shard_group, shards_configurations|
         next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
 
+        # configures path for migrations of this shard group and creates dir if necessary
         setup_migrations_path(shard_group)
 
-        shards_configurations.each do |shard, configuration|
+        shards_configurations.each do |shard, _|
           next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
           puts "== Loading schema of #{shard_group}:#{shard}"
 
@@ -125,7 +126,7 @@ shards_namespace = namespace :shards do
   end
 
   namespace :migrate do
-    desc  'Rollbacks the shards one migration and re migrate up (options: RAILS_ENV=x, VERSION=x, STEP=x, SHARD_GROUP=x, SHARD=x).'
+    desc 'Rollbacks the shards one migration and re migrate up (options: RAILS_ENV=x, VERSION=x, STEP=x, SHARD_GROUP=x, SHARD=x).'
     task redo: [:environment] do
       if ENV["VERSION"]
         shards_namespace["migrate:down"].invoke
@@ -141,15 +142,15 @@ shards_namespace = namespace :shards do
 
     desc 'Runs the "up" for a given migration VERSION.'
     task up: [:_make_activerecord_base_shardable] do
-      version = ENV["VERSION"] ? ENV["VERSION"].to_i : nil
-      raise "VERSION is required" unless version
+      version = get_version_or_else "VERSION is required"
 
       Rails::Sharding.configurations.each do |shard_group, shards_configurations|
         next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
 
+        # configures path for migrations of this shard group and creates dir if necessary
         setup_migrations_path(shard_group)
 
-        shards_configurations.each do |shard, configuration|
+        shards_configurations.each do |shard, _|
           next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
           puts "== Migrating up shard #{shard_group}:#{shard}"
           Rails::Sharding.using_shard(shard_group, shard) do
@@ -163,15 +164,15 @@ shards_namespace = namespace :shards do
 
     desc 'Runs the "down" for a given migration VERSION.'
     task down: [:_make_activerecord_base_shardable] do
-      version = ENV["VERSION"] ? ENV["VERSION"].to_i : nil
-      raise "VERSION is required - To go down one migration, run db:rollback" unless version
+      version = get_version_or_else "VERSION is required - To go down one migration, run db:rollback"
 
       Rails::Sharding.configurations.each do |shard_group, shards_configurations|
         next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
 
+        # configures path for migrations of this shard group and creates dir if necessary
         setup_migrations_path(shard_group)
 
-        shards_configurations.each do |shard, configuration|
+        shards_configurations.each do |shard, _|
           next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
           puts "== Migrating down shard #{shard_group}:#{shard}"
           Rails::Sharding.using_shard(shard_group, shard) do
@@ -190,9 +191,10 @@ shards_namespace = namespace :shards do
     Rails::Sharding.configurations.each do |shard_group, shards_configurations|
       next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
 
+      # configures path for migrations of this shard group and creates dir if necessary
       setup_migrations_path(shard_group)
 
-      shards_configurations.each do |shard, configuration|
+      shards_configurations.each do |shard, _|
         next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
         puts "== Rolling back shard #{shard_group}:#{shard}"
         Rails::Sharding.using_shard(shard_group, shard) do
@@ -206,15 +208,9 @@ shards_namespace = namespace :shards do
 
   desc "Retrieves the current schema version number"
   task version: [:_make_activerecord_base_shardable] do
-    Rails::Sharding.configurations.each do |shard_group, shards_configurations|
-      next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
-
-      shards_configurations.each do |shard, configuration|
-        next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
-
-        Rails::Sharding.using_shard(shard_group, shard) do
-          puts "Shard #{shard_group}:#{shard} version: #{ActiveRecord::Migrator.current_version}"
-        end
+    Rails::Sharding.for_each_shard(ENV["SHARD_GROUP"], ENV["SHARD"]) do |shard_group, shard, _configuration|
+      Rails::Sharding.using_shard(shard_group, shard) do
+        puts "Shard #{shard_group}:#{shard} version: #{ActiveRecord::Migrator.current_version}"
       end
     end
   end
@@ -226,9 +222,10 @@ shards_namespace = namespace :shards do
       Rails::Sharding.test_configurations.each do |shard_group, shards_configurations|
         next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
 
+        # configures path for migrations of this shard group and creates dir if necessary
         setup_migrations_path(shard_group)
 
-        shards_configurations.each do |shard, configuration|
+        shards_configurations.each do |shard, _|
           next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
 
           puts "== Loading test schema on shard #{shard_group}:#{shard}"
@@ -237,6 +234,10 @@ shards_namespace = namespace :shards do
             should_reconnect = Rails::Sharding::ConnectionHandler.connection_pool(shard_group, shard).active_connection?
             Rails::Sharding::ConnectionHandler.establish_connection(shard_group, shard, 'test')
 
+            # saves the current RAILS_ENV (we must change it so the environment is set correcly on the metadata table)
+            initial_rails_env = Rails.env
+            Rails.env = 'test'
+
             schema_filename = shard_schema_path(shard_group, shard)
             ActiveRecord::Tasks::DatabaseTasks.check_schema_file(schema_filename)
             Rails::Sharding.using_shard(shard_group, shard) do
@@ -244,6 +245,9 @@ shards_namespace = namespace :shards do
               load(schema_filename)
             end
           ensure
+            # restores rails env
+            Rails.env = initial_rails_env
+
             if should_reconnect
               # reestablishes connection for RAILS_ENV environment (whatever that is)
               Rails::Sharding::ConnectionHandler.establish_connection(shard_group, shard)
@@ -262,26 +266,20 @@ shards_namespace = namespace :shards do
 
     desc "Empty the test shards (drops all tables) (options: SHARD_GROUP=x, SHARD=x)"
     task :purge => [:_make_activerecord_base_shardable] do
-      Rails::Sharding.test_configurations.each do |shard_group, shards_configurations|
-        next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
-
-        shards_configurations.each do |shard, configuration|
-          next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
-
-          puts "== Purging test shard #{shard_group}:#{shard}"
-          begin
-            # establishes connection with test shard, saving if it was connected before (rails 4.2 doesn't do this, but should)
-            should_reconnect = Rails::Sharding::ConnectionHandler.connection_pool(shard_group, shard).active_connection?
-            Rails::Sharding::ConnectionHandler.establish_connection(shard_group, shard, 'test')
+      Rails::Sharding.for_each_shard(ENV["SHARD_GROUP"], ENV["SHARD"]) do |shard_group, shard, configuration|
+        puts "== Purging test shard #{shard_group}:#{shard}"
+        begin
+          # establishes connection with test shard, saving if it was connected before (rails 4.2 doesn't do this, but should)
+          should_reconnect = Rails::Sharding::ConnectionHandler.connection_pool(shard_group, shard).active_connection?
+          Rails::Sharding::ConnectionHandler.establish_connection(shard_group, shard, 'test')
 
-            Rails::Sharding.using_shard(shard_group, shard) do
-              ActiveRecord::Tasks::DatabaseTasks.purge(configuration)
-            end
-          ensure
-            if should_reconnect
-              # reestablishes connection for RAILS_ENV environment (whatever that is)
-              Rails::Sharding::ConnectionHandler.establish_connection(shard_group, shard)
-            end
+          Rails::Sharding.using_shard(shard_group, shard) do
+            ActiveRecord::Tasks::DatabaseTasks.purge(configuration)
+          end
+        ensure
+          if should_reconnect
+            # reestablishes connection for RAILS_ENV environment (whatever that is)
+            Rails::Sharding::ConnectionHandler.establish_connection(shard_group, shard)
           end
         end
       end
@@ -304,4 +302,10 @@ shards_namespace = namespace :shards do
     FileUtils.mkdir_p(shard_group_schemas_dir)
     File.join(shard_group_schemas_dir, shard_name + "_schema.rb")
   end
+
+  def get_version_or_else(error_message='VERSION is required')
+    version = ENV["VERSION"] ? ENV["VERSION"].to_i : nil
+    raise error_message unless version
+    version
+  end
 end

data/rails-sharding.gemspec

@@ -23,12 +23,12 @@ Gem::Specification.new do |spec|
   spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
   spec.require_paths = ["lib"]
 
-  spec.add_runtime_dependency 'rails', '>= 4.2', '< 5.0'
+  spec.add_runtime_dependency 'rails', '~> 5.0'
 
   spec.add_development_dependency "bundler", "~> 1.12"
-  spec.add_development_dependency "rake", "~> 10.0"
+  spec.add_development_dependency "rake", "~> 11.0"
   spec.add_development_dependency "rspec", "~> 3.0"
-  spec.add_development_dependency "byebug", '~> 0'
+  spec.add_development_dependency "byebug", '~> 9'
   spec.add_development_dependency "mysql2", '~> 0'
   spec.add_development_dependency "codeclimate-test-reporter", '~> 0'
 end

metadata

@@ -1,33 +1,27 @@
 --- !ruby/object:Gem::Specification
 name: rails-sharding
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 1.0.1
 platform: ruby
 authors:
 - Henrique Gubert
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2016-10-07 00:00:00.000000000 Z
+date: 2016-11-02 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rails
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '4.2'
-    - - "<"
+    - - "~>"
       - !ruby/object:Gem::Version
         version: '5.0'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '4.2'
-    - - "<"
+    - - "~>"
       - !ruby/object:Gem::Version
         version: '5.0'
 - !ruby/object:Gem::Dependency
@@ -50,14 +44,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '10.0'
+        version: '11.0'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '10.0'
+        version: '11.0'
 - !ruby/object:Gem::Dependency
   name: rspec
   requirement: !ruby/object:Gem::Requirement
@@ -78,14 +72,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '9'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '9'
 - !ruby/object:Gem::Dependency
   name: mysql2
   requirement: !ruby/object:Gem::Requirement