rails-sharding 0.1.1 → 1.0.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 05fd12669243e14c2d6f1b1748acf461b8316dbb
-  data.tar.gz: d1508fddee17b0a5c7eb0203b8dd1c9b52b74c66
+  metadata.gz: cbec123e1537bac8bf7f3b6d2e711cfefa156f41
+  data.tar.gz: 4fae6ab76f8b687746436585a74d618c5f4a2308
 SHA512:
-  metadata.gz: 5ea75e62463c090f6806baedfd648768007c192d4e3a9e162c162083eb015cacb8eb6dfbc8ed0ce739bb9cc17c9a3889abc6b0bd02bb24f3f76470478996e85c
-  data.tar.gz: 524dfe1e8c7ea123209746accfdb726ea0a6011f43c460cbcf2cff3c11531b2ee7afa99c26091d54ad8792dd05dca62aff2badaee5290babb1f72631ed0e94d0
+  metadata.gz: f8488d56f2e5eb38df46f138d4ad7ac28bc12008ee08383f984268fd3201629bb3eebeb54c96120bddf846e2e1f39f171e38fc4ab52457bbcdbcafa805b0e72c
+  data.tar.gz: 95a78672f724efa645c911d0d824b57be1dcf43d96257be54aab15e92ba544b81ce370e3a9131ee5abdd27f4605a8a49a5e45dd66dd00f4e4f58b3c7dfc48f28

data/README.md

@@ -27,7 +27,8 @@ 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 0.1.1 -> compatible with Rails 4.2
+Gem version 1.0.0 -> compatible with Rails 5.0
 
 ## Installation
 
@@ -83,7 +84,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
@@ -154,4 +155,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/lib/rails/sharding/connection_handler.rb

@@ -17,40 +17,43 @@ 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_handler.retrieve_connection(connection_name(shard_group, shard_name))
     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)
@@ -58,7 +61,7 @@ module Rails::Sharding
     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,35 +73,11 @@ 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"
     def self.connection_name(shard_group, shard_name)
       shard_group.to_s + ':' + shard_name.to_s
     end
-
-    class ConnectionPoolOwner
-      attr_reader :name
-
-      def initialize(name)
-      	@name = name
-      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
-
-      # in case owner ends up printed by rails in an error message when retrieving connection
-      def to_s
-        "ConnectionPoolOwner with name #{self.name}"
-      end
-    end
   end
 end

data/lib/rails/sharding/core.rb

@@ -14,8 +14,14 @@ module Rails::Sharding
 
       ShardThreadRegistry.current_shard_group = shard_group
       ShardThreadRegistry.current_shard_name = shard_name
+      ShardThreadRegistry.shard_connection_used = false
       yield
     ensure
+      # shows warning to user
+      if !ShardThreadRegistry.shard_connection_used
+        puts "Warning: no connection to shard '#{ShardThreadRegistry.current_shard_group}:#{ShardThreadRegistry.current_shard_name}' was made inside the using_shard block. Make sure you don't forget to include Rails::Sharding::ShardableModel to the models you want to be sharded"
+      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
@@ -26,7 +32,7 @@ module Rails::Sharding
     def self.configurations(environment=Rails.env)
       @@db_configs ||= YAML.load_file(Config.shards_config_file)
       @@db_configs[environment]
-    rescue Errno::ENOENT => e
+    rescue Errno::ENOENT
       raise Errors::ConfigNotFoundError, Config.shards_config_file.to_s + ' file was not found'
     end
 
@@ -46,6 +52,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,13 @@
-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 variables to store the current shard of connection
+    thread_mattr_accessor :_current_shard_group
+    thread_mattr_accessor :_current_shard_name
+
+    # auxiliary variable used to check if shard connectio was used inside an
+    # using_shard block (so we can print an alert if not)
+    thread_mattr_accessor :shard_connection_used
 
     def self.connecting_to_master?
       current_shard_group.nil? || current_shard_name.nil?
@@ -18,6 +20,7 @@ module Rails::Sharding
     def self.connect_back_to_master!
       self.current_shard_group = nil
       self.current_shard_name = nil
+      self.shard_connection_used = false
     end
 
     # Returns the current shard group (for the current Thread)

data/lib/rails/sharding/shardable_model.rb

@@ -35,6 +35,7 @@ module Rails::Sharding
         if ShardThreadRegistry.connecting_to_master?
           return original_connection_pool
         else
+          ShardThreadRegistry.shard_connection_used = true # records that shard connection was used at least once
           return ConnectionHandler.connection_pool(*ShardThreadRegistry.current_shard_group_and_name)
         end
       end
@@ -44,6 +45,7 @@ module Rails::Sharding
         if ShardThreadRegistry.connecting_to_master?
           return original_retrieve_connection
         else
+          ShardThreadRegistry.shard_connection_used = true # records that shard connection was used at least once
           return ConnectionHandler.retrieve_connection(*ShardThreadRegistry.current_shard_group_and_name)
         end
       end

data/lib/rails/sharding/version.rb

@@ -1,5 +1,5 @@
 module Rails
   module Sharding
-    VERSION = "0.1.1"
+    VERSION = "1.0.0"
   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,7 +23,7 @@ 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"

metadata

@@ -1,33 +1,27 @@
 --- !ruby/object:Gem::Specification
 name: rails-sharding
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 1.0.0
 platform: ruby
 authors:
 - Henrique Gubert
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2016-10-07 00:00:00.000000000 Z
+date: 2016-10-10 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