lib/octopus/migration.rb in ar-octopus-0.9.1 vs lib/octopus/migration.rb in ar-octopus-0.9.2

@@ -75,10 +75,13 @@
           alias_method :down_without_octopus, :down
           alias_method :down, :down_with_octopus
 
           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
 
       base.send :alias_method, :run_without_octopus, :run
       base.send :alias_method, :run, :run_with_octopus
@@ -140,9 +143,17 @@
       def run_with_octopus(direction, migrations_paths, target_version)
         return run_without_octopus(direction, migrations_paths, target_version) unless connection.is_a?(Octopus::Proxy)
 
         connection.send_queries_to_multiple_shards(connection.shard_names) do
           run_without_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