lib/tasks/switchman.rake in switchman-1.5.9 vs lib/tasks/switchman.rake in switchman-1.5.10

@@ -47,11 +47,11 @@
     def self.shardify_task(task_name)
       old_task = ::Rake::Task[task_name]
       old_actions = old_task.actions.dup
       old_task.actions.clear
 
-      old_task.enhance do
+      old_task.enhance do |*task_args|
         if ::Rails.env.test?
           require 'switchman/test_helper'
           TestHelper.recreate_persistent_test_shards(dont_create: true)
         end
 
@@ -60,10 +60,10 @@
             shard = Shard.current
             puts "#{shard.id}: #{shard.description}"
             ::ActiveRecord::Base.connection_pool.spec.config[:shard_name] = Shard.current.name
             ::ActiveRecord::Base.configurations[::Rails.env] = ::ActiveRecord::Base.connection_pool.spec.config.stringify_keys
             shard.database_server.unshackle do
-              old_actions.each(&:call)
+              old_actions.each { |action| action.call(*task_args) }
             end
             nil
           end
         end
       end