lib/octopus/migration.rb in ar-octopus-0.0.24 vs lib/octopus/migration.rb in ar-octopus-0.0.25

@@ -5,20 +5,22 @@
         announce_without_octopus("#{message} - #{get_current_shard}")
       end
       
       alias_method_chain :migrate, :octopus
       alias_method_chain :announce, :octopus
+      attr_accessor :current_shard
     end
   end
 
   def using(*args, &block)
     if self.connection().is_a?(Octopus::Proxy) && !block_given?
       args.each do |shard|
         self.connection().check_schema_migrations(shard)
       end
 
       self.connection().block = true
+      self.current_shard = args
       self.connection().current_shard = args        
     end
     
     if block_given?
       self.connection.run_queries_on_shard(args, &block)
@@ -49,14 +51,15 @@
   end
 
 
   def migrate_with_octopus(direction)
     conn = ActiveRecord::Base.connection
+    return migrate_without_octopus(direction) unless conn.is_a?(Octopus::Proxy)
+    self.connection().current_shard = self.current_shard if self.current_shard != nil
+    
     groups = conn.instance_variable_get(:@groups)
     
     begin
-      return migrate_without_octopus(direction) unless conn.is_a?(Octopus::Proxy)
-    
       if conn.current_group.is_a?(Array)
         conn.current_group.each { |group| conn.send_queries_to_multiple_shards(groups[group]) { migrate_without_octopus(direction) } } 
       elsif conn.current_group.is_a?(Symbol)       
         conn.send_queries_to_multiple_shards(groups[conn.current_group]) { migrate_without_octopus(direction) }     
       elsif conn.current_shard.is_a?(Array)