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

@@ -1,37 +1,35 @@
 module Octopus::Migration  
   def self.extended(base)
     class << base
-      alias_method_chain :migrate, :octopus
-      
-      def announce(message)
-        version = defined?(@version) ? @version : nil
-
-        text = "#{version} #{name}: #{message} - #{get_current_shard}"
-        length = [0, 75 - text.length].max
-        write "== %s %s" % [text, "=" * length]
+      def announce_with_octopus(message)
+        announce_without_octopus("#{message} - #{get_current_shard}")
       end
+      
+      alias_method_chain :migrate, :octopus
+      alias_method_chain :announce, :octopus
     end
   end
 
   def using(*args, &block)
-    if self.connection().is_a?(Octopus::Proxy)
+    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.connection().current_shard = args        
     end
     
-    yield if block_given?
+    if block_given?
+      self.connection.run_queries_on_shard(args, &block)
+    end 
 
     return self
   end
 
   def using_group(*args)
-    
     if self.connection().is_a?(Octopus::Proxy)
       args.each do |group_shard|
         shards = self.connection().instance_variable_get(:@groups)[group_shard] || []
 
         shards.each do |shard|
@@ -53,22 +51,24 @@
 
   def migrate_with_octopus(direction)
     conn = ActiveRecord::Base.connection
     groups = conn.instance_variable_get(:@groups)
     
-    return migrate_without_octopus(direction) unless conn.is_a?(Octopus::Proxy)
+    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)
-      conn.send_queries_to_multiple_shards(conn.current_shard) { migrate_without_octopus(direction) }     
-    else
-      migrate_without_octopus(direction)
+      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)
+        conn.send_queries_to_multiple_shards(conn.current_shard) { migrate_without_octopus(direction) }     
+      else
+        migrate_without_octopus(direction)
+      end
+    ensure
+      conn.clean_proxy()
     end
-
-    conn.clean_proxy()
   end
 end
 
 ActiveRecord::Migration.extend(Octopus::Migration)