switchman.rake in switchman-1.6.0

- old
+ new

@@ -42,11 +42,15 @@
 
     def self.options
       { parallel: ENV['PARALLEL'].to_i, max_procs: ENV['MAX_PARALLEL_PROCS'] }
     end
 
-    def self.shardify_task(task_name)
+    # categories - an array or proc, to activate as the current shard during the
+    # task. tasks which modify the schema may want to pass all categories in
+    # so that schema updates for non-default tables happen against all shards.
+    # this is handled automatically for the default migration tasks, below.
+    def self.shardify_task(task_name, categories: [:default])
       old_task = ::Rake::Task[task_name]
       old_actions = old_task.actions.dup
       old_task.actions.clear
 
       old_task.enhance do |*task_args|
@@ -55,11 +59,12 @@
           TestHelper.recreate_persistent_test_shards(dont_create: true)
         end
 
         ::Shackles.activate(:deploy) do
           begin
-            Shard.with_each_shard(scope, Shard.categories, options) do
+            categories = categories.call if categories.respond_to?(:call)
+            Shard.with_each_shard(scope, categories, options) do
               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
@@ -73,10 +78,12 @@
           end
         end
       end
     end
 
-    %w{db:migrate db:migrate:up db:migrate:down db:rollback}.each { |task_name| shardify_task(task_name) }
+    %w{db:migrate db:migrate:up db:migrate:down db:rollback}.each do |task_name|
+      shardify_task(task_name, categories: ->{ Shard.categories })
+    end
 
     private
 
     def self.shard_scope(scope, raw_shard_ids)
       raw_shard_ids = raw_shard_ids.split(',')