rails-sharding 0.1.0

data/lib/rails/sharding/config.rb

@@ -0,0 +1,32 @@
+module Rails::Sharding
+  class Config
+    DEFAULT_CONFIGS = {
+      # If true one connection will be established per shard (in every shard group)
+      # on startup. This only establishes the connection with the database but
+      # it does not retrieve a connection yet. This will be done by the ConnectionPool
+      # when necessary.
+      # If false the user must call Shards::ConnectionHandler.establish_connection(shard_group, shard_name)
+      # manually at least once before using each shard.
+      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
+      extend_active_record_scope: true,
+
+      # Specifies where to find the definition of the shards configurations
+      shards_config_file: 'config/shards.yml',
+
+      # Specifies where to find the migrations for each shard group
+      shards_migrations_dir: 'db/shards_migrations',
+
+      # Specifies where to find the schemas for each shard group
+      shards_schemas_dir: 'db/shards_schemas'
+    }
+
+    DEFAULT_CONFIGS.each do |config_name, default_value|
+      self.cattr_accessor config_name
+      self.send(config_name.to_s + '=', default_value)
+    end
+
+  end
+end

data/lib/rails/sharding/connection_handler.rb

@@ -0,0 +1,104 @@
+module Rails::Sharding
+  class ConnectionHandler
+
+    # Establishes connections to all shards in all shard groups.
+    # Despite the name, this actually only creates a connection pool with zero
+    # connections for each shard. The connections will be allocated for each
+    # thread when #retrieve_connection or #with_connection are called
+    def self.establish_all_connections
+      Core.shard_groups.each do |shard_group|
+        Core.shard_names(shard_group).each do |shard_name|
+          establish_connection(shard_group, shard_name)
+        end
+      end
+    end
+
+    # Establishes a connection to a single shard in a single shard group
+    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))
+        raise Errors::ConfigNotFoundError, "Cannot find configuration for environment '#{environment}' in #{Config.shards_config_file}"
+      end
+
+      unless shard_group_configurations = configurations[shard_group.to_s]
+        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
+        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)
+    end
+
+    def self.connection_pool(shard_group, shard_name)
+      connection_handler.retrieve_connection_pool(connection_pool_owner(shard_group, shard_name))
+    rescue Errors::ConnectionPoolRetrievalError
+      # 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)}"
+    end
+
+    def self.retrieve_connection(shard_group, shard_name)
+      connection_handler.retrieve_connection(connection_pool_owner(shard_group, shard_name))
+    end
+
+    def self.connected?(shard_group, shard_name)
+      connection_handler.connected?(connection_pool_owner(shard_group, shard_name))
+    end
+
+    def self.with_connection(shard_group, shard_name, &block)
+    	connection_pool(shard_group, shard_name).with_connection(&block)
+    end
+
+    def self.remove_connection(shard_group, shard_name)
+      connection_handler.remove_connection(connection_pool_owner(shard_group, shard_name))
+    end
+
+  private
+
+    def self.connection_handler
+      raise Errors::UninitializedError, 'Shards::ConnectionHandler was not setup' unless defined? @@connection_handler
+      @@connection_handler
+    end
+
+    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

@@ -0,0 +1,67 @@
+require 'rails/sharding/active_record_extensions'
+require 'rails/sharding/config'
+require 'rails/sharding/connection_handler'
+require 'rails/sharding/errors'
+require 'rails/sharding/shard_thread_registry'
+require 'rails/sharding/shardable_model'
+
+module Rails::Sharding
+  class Core
+
+    # 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
+      yield
+    ensure
+      # 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!
+    end
+
+    def self.configurations(environment=Rails.env)
+      @@db_configs ||= YAML.load_file(Config.shards_config_file)
+      @@db_configs[environment]
+    rescue Errno::ENOENT => e
+      raise Errors::ConfigNotFoundError, Config.shards_config_file.to_s + ' file was not found'
+    end
+
+    def self.reset_configurations_cache
+      @@db_configs = nil
+    end
+
+    def self.test_configurations
+      self.configurations('test')
+    end
+
+    def self.shard_groups
+      self.configurations.keys
+    end
+
+    def self.shard_names(shard_group)
+      self.configurations[shard_group.to_s].keys
+    end
+
+    # Method that should be called on a rails initializer
+    def self.setup
+      if block_given?
+        yield Config
+      end
+
+      if Config.establish_all_connections_on_setup
+        # Establishes connections with all shards specified in config/shards.yml
+        ConnectionHandler.establish_all_connections
+      end
+
+      if Config.extend_active_record_scope
+        # includes the #using_shard method to all AR scopes
+        ActiveRecordExtensions.extend_active_record_scope
+      end
+    end
+
+  end
+end

data/lib/rails/sharding/errors.rb

@@ -0,0 +1,13 @@
+module Rails::Sharding
+  module Errors
+
+    class UninitializedError < StandardError; end
+
+    class WrongUsageError < StandardError; end
+
+    class ConnectionPoolRetrievalError < StandardError; end
+
+    class ConfigNotFoundError < StandardError; end
+
+  end
+end

data/lib/rails/sharding/railtie.rb

@@ -0,0 +1,10 @@
+
+module Rails::Sharding
+
+  class Railtie < Rails::Railtie
+    rake_tasks do
+      load "tasks/rails-sharding.rake"
+    end
+  end
+
+end

data/lib/rails/sharding/shard_thread_registry.rb

@@ -0,0 +1,48 @@
+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
+
+    def self.connecting_to_master?
+      current_shard_group.nil? || current_shard_name.nil?
+    end
+
+    def self.connecting_to_shard?
+      !connecting_to_master?
+    end
+
+    def self.connect_back_to_master!
+      self.current_shard_group = nil
+      self.current_shard_name = nil
+    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
+    end
+
+    # Returns the current shard name (for the current Thread)
+    def self.current_shard_name
+      _current_shard_name
+    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
+    end
+
+    def self.current_shard_group_and_name
+      [current_shard_group, current_shard_name]
+    end
+
+  end
+end

data/lib/rails/sharding/shardable_model.rb

@@ -0,0 +1,109 @@
+# The purpose of this class is to override ActiveRecord::ConnectionHandling
+# methods that access the DB connection or connection pool in any way.
+#
+# If Rails::Sharding::ShardThreadRegistry.connecting_to_master? is true, we just
+# delegate all calls to the original methods of ActiveRecord::ConnectionHandling,
+# otherwise we access our own ConnectionHandler and retrieve the connection
+# or connection pool to the selected shard
+
+module Rails::Sharding
+  module ShardableModel
+
+    def self.included(base)
+      # base.extend ClassMethods
+      base.extend ClassMethodsOverrides
+    end
+
+    # Module that includes all class methods that will be overriden on the model
+    # class when Rails::Sharding::ShardableModel is included
+    module ClassMethodsOverrides
+      # dinamically saves original methods with prefix "original_" and overrides
+      # then with the methods with prefix "sharded_"
+      def self.extended(klass)
+        self.instance_methods.each do |sharded_method_name|
+          method_name = sharded_method_name.to_s.match(/sharded_(.+)/)[1].to_sym
+
+          klass.singleton_class.instance_eval do
+            alias_method "original_#{method_name}".to_sym, method_name
+            alias_method method_name, sharded_method_name
+          end
+        end
+      end
+
+      # @overrides ActiveRecord::ConnectionHandling#connection_pool
+      def sharded_connection_pool
+        if ShardThreadRegistry.connecting_to_master?
+          return original_connection_pool
+        else
+          return ConnectionHandler.connection_pool(*ShardThreadRegistry.current_shard_group_and_name)
+        end
+      end
+
+      # @overrides ActiveRecord::ConnectionHandling#retrieve_connection
+      def sharded_retrieve_connection
+        if ShardThreadRegistry.connecting_to_master?
+          return original_retrieve_connection
+        else
+          return ConnectionHandler.retrieve_connection(*ShardThreadRegistry.current_shard_group_and_name)
+        end
+      end
+
+      # @overrides ActiveRecord::ConnectionHandling#sharded_connected?
+      def sharded_connected?
+        if ShardThreadRegistry.connecting_to_master?
+          return original_connected?
+        else
+          return ConnectionHandler.connected?(*ShardThreadRegistry.current_shard_group_and_name)
+        end
+      end
+
+      # @overrides ActiveRecord::ConnectionHandling#remove_connection (only if no parameters are passed)
+      def sharded_remove_connection(klass=nil)
+        if ShardThreadRegistry.connecting_to_master? || klass
+          return klass ? original_remove_connection(klass) : original_remove_connection
+        else
+          return ConnectionHandler.remove_connection(*ShardThreadRegistry.current_shard_group_and_name)
+        end
+      end
+
+      # @overrides ActiveRecord::ConnectionHandling#establish_connection
+      # In the case we are connecting to a shard, ignore spec parameter and use
+      # what is in ShardThreadRegistry instead
+      def sharded_establish_connection(spec=nil)
+        if ShardThreadRegistry.connecting_to_master?
+          return original_establish_connection(spec)
+        else
+          return ConnectionHandler.establish_connection(*ShardThreadRegistry.current_shard_group_and_name)
+        end
+      end
+
+      # @overrides ActiveRecord::ConnectionHandling#clear_active_connections!
+      def sharded_clear_active_connections!
+        if ShardThreadRegistry.connecting_to_master?
+          return original_clear_active_connections!
+        else
+          return ConnectionHandler.connection_handler.clear_active_connections!
+        end
+      end
+
+      # @overrides ActiveRecord::ConnectionHandling#clear_reloadable_connections!
+      def sharded_clear_reloadable_connections!
+        if ShardThreadRegistry.connecting_to_master?
+          return original_clear_reloadable_connections!
+        else
+          return ConnectionHandler.connection_handler.clear_reloadable_connections!
+        end
+      end
+
+      # @overrides ActiveRecord::ConnectionHandling#clear_all_connections!
+      def sharded_clear_all_connections!
+        if ShardThreadRegistry.connecting_to_master?
+          return original_clear_all_connections!
+        else
+          return ConnectionHandler.connection_handler.clear_all_connections!
+        end
+      end
+    end
+
+  end
+end

data/lib/rails/sharding/version.rb

@@ -0,0 +1,5 @@
+module Rails
+  module Sharding
+    VERSION = "0.1.0"
+  end
+end

data/lib/tasks/rails-sharding.rake

@@ -0,0 +1,309 @@
+require "active_record"
+
+shards_namespace = namespace :shards do
+  task _make_activerecord_base_shardable: [:environment] do
+    # Several resources used (like Migrator, SchemaDumper, schema methods)
+    # implicitly use ActiveRecord::Base.connection, so we have to make it
+    # shardable so we can call using_shard and switch the connection
+    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
+    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
+        puts "== Creating shard #{shard_group}:#{shard}"
+        ActiveRecord::Tasks::DatabaseTasks.create(configuration)
+      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
+
+      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
+    end
+  end
+
+  desc "Migrate the database (options: RAILS_ENV=x, SHARD_GROUP=x, VERSION=x, VERBOSE=false, SCOPE=blog)."
+  task migrate: [:_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
+
+      # 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)
+
+      shards_configurations.each do |shard, configuration|
+        next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
+        puts "== Migrating shard #{shard_group}:#{shard}"
+        Rails::Sharding.using_shard(shard_group, shard) do
+          ActiveRecord::Tasks::DatabaseTasks.migrate
+        end
+      end
+    end
+
+    shards_namespace["_dump"].invoke
+  end
+
+  # IMPORTANT: This task won't dump the schema if ActiveRecord::Base.dump_schema_after_migration is set to false
+  task :_dump do
+    if ActiveRecord::Base.dump_schema_after_migration
+      case ActiveRecord::Base.schema_format
+      when :ruby
+        shards_namespace["schema:dump"].invoke
+      when :sql
+        raise "sql schema dump not supported by shards"
+      else
+        raise "unknown schema format #{ActiveRecord::Base.schema_format}"
+      end
+    end
+    # Allow this task to be called as many times as required. An example is the
+    # migrate:redo task, which calls other two internally that depend on this one.
+    shards_namespace["_dump"].reenable
+  end
+
+  namespace :schema do
+    desc "Creates a schema.rb for each shard that is portable against any DB supported by Active Record (options: RAILS_ENV=x, SHARD_GROUP=x, SHARD=x)"
+    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
+
+        # configures path for schemas of this shard group and creates dir if necessary
+        shard_group_schemas_dir = File.join(Rails::Sharding::Config.shards_schemas_dir, shard_group.to_s)
+        FileUtils.mkdir_p(shard_group_schemas_dir)
+
+        shards_configurations.each do |shard, configuration|
+          next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
+          puts "== Dumping schema of #{shard_group}:#{shard}"
+
+          schema_filename = File.join(shard_group_schemas_dir, shard + "_schema.rb")
+          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
+      end
+
+      # Allow this task to be called as many times as required. An example is the
+      # migrate:redo task, which calls other two internally that depend on this one.
+      shards_namespace["schema:dump"].reenable
+    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
+      Rails::Sharding.configurations.each do |shard_group, shards_configurations|
+        next if ENV["SHARD_GROUP"] && ENV["SHARD_GROUP"] != shard_group.to_s
+
+        # configures path for schemas of this shard group
+        shard_group_schemas_path = 'db/shards_schemas/' + shard_group.to_s
+        Rails.application.paths.add shard_group_schemas_path
+        shard_group_schemas_dir = Rails.application.paths[shard_group_schemas_path].to_a.first
+
+        shards_configurations.each do |shard, configuration|
+          next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
+          puts "== Loading schema of #{shard_group}:#{shard}"
+
+          schema_filename = File.join(shard_group_schemas_dir, shard + "_schema.rb")
+          ActiveRecord::Tasks::DatabaseTasks.check_schema_file(schema_filename)
+          Rails::Sharding.using_shard(shard_group, shard) do
+            load(schema_filename)
+          end
+        end
+      end
+    end
+
+    task load_if_ruby: ["shards:create", :environment] do
+      shards_namespace["schema:load"].invoke if ActiveRecord::Base.schema_format == :ruby
+    end
+  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).'
+    task redo: [:environment] do
+      if ENV["VERSION"]
+        shards_namespace["migrate:down"].invoke
+        shards_namespace["migrate:up"].invoke
+      else
+        shards_namespace["rollback"].invoke
+        shards_namespace["migrate"].invoke
+      end
+    end
+
+    desc 'Resets your shards using your migrations for the current environment'
+    task reset: ["shards:drop", "shards:create", "shards:migrate"]
+
+    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
+
+      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
+        shard_group_migrations_path = 'db/shards_migrations/' + shard_group.to_s
+        Rails.application.paths.add shard_group_migrations_path
+        ActiveRecord::Tasks::DatabaseTasks.migrations_paths = Rails.application.paths[shard_group_migrations_path].to_a
+
+        shards_configurations.each do |shard, configuration|
+          next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
+          puts "== Migrating up shard #{shard_group}:#{shard}"
+          Rails::Sharding.using_shard(shard_group, shard) do
+            ActiveRecord::Migrator.run(:up, ActiveRecord::Tasks::DatabaseTasks.migrations_paths, version)
+          end
+        end
+      end
+
+      shards_namespace["_dump"].invoke
+    end
+
+    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
+
+      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
+        shard_group_migrations_path = 'db/shards_migrations/' + shard_group.to_s
+        Rails.application.paths.add shard_group_migrations_path
+        ActiveRecord::Tasks::DatabaseTasks.migrations_paths = Rails.application.paths[shard_group_migrations_path].to_a
+
+        shards_configurations.each do |shard, configuration|
+          next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
+          puts "== Migrating down shard #{shard_group}:#{shard}"
+          Rails::Sharding.using_shard(shard_group, shard) do
+            ActiveRecord::Migrator.run(:down, ActiveRecord::Tasks::DatabaseTasks.migrations_paths, version)
+          end
+        end
+      end
+
+      shards_namespace["_dump"].invoke
+    end
+  end
+
+  desc "Rolls the schema back to the previous version (options: RAILS_ENV=x, STEP=x, SHARD_GROUP=x, SHARD=x)."
+  task rollback: [:_make_activerecord_base_shardable] do
+    step = ENV["STEP"] ? ENV["STEP"].to_i : 1
+    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
+      shard_group_migrations_path = 'db/shards_migrations/' + shard_group.to_s
+      Rails.application.paths.add shard_group_migrations_path
+      ActiveRecord::Tasks::DatabaseTasks.migrations_paths = Rails.application.paths[shard_group_migrations_path].to_a
+
+      shards_configurations.each do |shard, configuration|
+        next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
+        puts "== Rolling back shard #{shard_group}:#{shard}"
+        Rails::Sharding.using_shard(shard_group, shard) do
+          ActiveRecord::Migrator.rollback(ActiveRecord::Tasks::DatabaseTasks.migrations_paths, step)
+        end
+      end
+    end
+
+    shards_namespace["_dump"].invoke
+  end
+
+  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
+      end
+    end
+  end
+
+
+  namespace :test do
+    desc "Recreate the test shards from existent schema files (options: SHARD_GROUP=x, SHARD=x)"
+    task load_schema: ['shards:test:purge'] 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 schemas of this shard group
+        shard_group_schemas_path = 'db/shards_schemas/' + shard_group.to_s
+        Rails.application.paths.add shard_group_schemas_path
+        shard_group_schemas_dir = Rails.application.paths[shard_group_schemas_path].to_a.first
+
+        shards_configurations.each do |shard, configuration|
+          next if ENV["SHARD"] && ENV["SHARD"] != shard.to_s
+
+          puts "== Loading test schema on shard #{shard_group}:#{shard}"
+          begin
+            # establishes connection with test shard, saving if it was connected before
+            should_reconnect = Rails::Sharding::ConnectionHandler.connection_pool(shard_group, shard).active_connection?
+            Rails::Sharding::ConnectionHandler.establish_connection(shard_group, shard, 'test')
+
+            schema_filename = File.join(shard_group_schemas_dir, shard + "_schema.rb")
+            ActiveRecord::Tasks::DatabaseTasks.check_schema_file(schema_filename)
+            Rails::Sharding.using_shard(shard_group, shard) do
+              ActiveRecord::Schema.verbose = false
+              load(schema_filename)
+            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
+      end
+    end
+
+    desc 'Load the test schema into the shards (options: SHARD_GROUP=x, SHARD=x)'
+    task prepare: [:environment] do
+      unless Rails::Sharding.test_configurations.blank?
+        shards_namespace["test:load_schema"].invoke
+      end
+    end
+
+    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.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
+      end
+    end
+  end
+end