Fixing issue described in GH thiagopradi#407

lib/octopus/migration.rb

@@ -77,6 +77,9 @@ class << self
 
           alias_method :run_without_octopus, :run
           alias_method :run, :run_with_octopus
+
+          alias_method :rollback_without_octopus, :rollback
+          alias_method :rollback, :rollback_with_octopus
         end
       end
 
@@ -145,6 +148,14 @@ def run_with_octopus(direction, migrations_paths, target_version)
         end
       end
 
+      def rollback_with_octopus(migrations_paths, steps = 1)
+        return rollback_without_octopus(migrations_paths, steps) unless connection.is_a?(Octopus::Proxy)
+
+        connection.send_queries_to_multiple_shards(connection.shard_names) do
+          rollback_without_octopus(migrations_paths, steps)
+        end
+      end
+
       private
 
       def connection

spec/octopus/migration_spec.rb

@@ -30,6 +30,25 @@
     end
   end
 
+  it "should rollback correctly migrations" do
+    migrations_root = File.expand_path(File.join(File.dirname(__FILE__), '..', 'migrations'))
+
+    ActiveRecord::Migrator.run(:up, migrations_root, 4)
+
+    expect(User.using(:canada).find_by_name('Group')).not_to be_nil
+    expect(User.using(:brazil).find_by_name('Group')).not_to be_nil
+    expect(User.using(:russia).find_by_name('Group')).not_to be_nil
+
+
+    Octopus.using(:canada) do
+      ActiveRecord::Migrator.rollback(migrations_root, 4)
+    end
+
+    expect(User.using(:canada).find_by_name('Group')).to be_nil
+    expect(User.using(:brazil).find_by_name('Group')).to be_nil
+    expect(User.using(:russia).find_by_name('Group')).to be_nil
+  end
+
   it 'should run once per shard' do
     OctopusHelper.migrating_to_version 5 do
       expect(User.using(:canada).where(:name => 'MultipleGroup').size).to eq(1)