Allow transform-tree to work on more than 1 tree at once.

Refactor various transform commands to remove cut and pasting.
twitter-archive · Nov 30, 2011 · 875190c · 875190c
1 parent 9733684
commit 875190c
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 84 deletions.
diff --git a/lib/gizzard/commands.rb b/lib/gizzard/commands.rb
@@ -744,32 +744,29 @@ def run
     end
   end
 
-  class TransformTreeCommand < Command
+  class BaseTransformCommand < Command
     def run
-      help!("wrong number of arguments") unless @argv.length == 2
-
       scheduler_options = command_options.scheduler_options || {}
-      template_s, shard_id_s = @argv
-
-      to_template    = ShardTemplate.parse(template_s)
-      shard_id       = ShardId.parse(shard_id_s)
-      base_name      = shard_id.table_prefix.split('_').first
-      forwarding     = manager.get_forwarding_for_shard(shard_id)
-      manifest       = manager.manifest(forwarding.table_id)
-      shard          = manifest.trees[forwarding]
-      copy_wrapper   = scheduler_options[:copy_wrapper]
-      be_quiet       = global_options.force && command_options.quiet
-      transformation = Transformation.new(shard.template, to_template, copy_wrapper)
+      be_quiet          = global_options.force && command_options.quiet
 
       scheduler_options[:quiet] = be_quiet
 
-      if transformation.noop?
+      transformations = get_transformations
+      transformations.reject! {|t,trees| t.noop? or trees.empty? }
+
+      if transformations.empty?
         puts "Nothing to do!"
         exit
       end
 
+      base_name = transformations.values.find {|v| v.is_a?(Hash) && !v.values.empty? }.values.find {|v| !v.nil?}.id.table_prefix.split('_').first
+
       unless be_quiet
-        puts transformation.inspect
+        transformations.each do |transformation, trees|
+          puts transformation.inspect
+          puts "Applied to #{trees.length} shards"
+          #trees.keys.sort.each {|f| puts "  #{f.inspect}" }
+        end
         puts ""
       end
 
@@ -781,24 +778,53 @@ def run
 
       Gizzard.schedule! manager,
                         base_name,
-                        { transformation => { forwarding => shard } },
+                        transformations,
                         scheduler_options
     end
   end
 
-  class TransformCommand < Command
-    def run
+  class TransformTreeCommand < BaseTransformCommand
+    def get_transformations
+      help!("must have an even number of arguments") unless @argv.length % 2 == 0
+
+      scheduler_options = command_options.scheduler_options || {}
+      copy_wrapper   = scheduler_options[:copy_wrapper]
+      skip_copies    = scheduler_options[:skip_copies] || false
+      transformations   = {}
+
+      memoized_transforms = {}
+      @argv.each_slice(2) do |(template_s, shard_id_s)|
+        to_template    = ShardTemplate.parse(template_s)
+        shard_id       = ShardId.parse(shard_id_s)
+        base_name      = shard_id.table_prefix.split('_').first
+        forwarding     = manager.get_forwarding_for_shard(shard_id)
+        manifest       = manager.manifest(forwarding.table_id)
+        shard          = manifest.trees[forwarding]
+
+        transform_args = [shard.template, to_template, copy_wrapper, skip_copies]
+        transformation = memoized_transforms.fetch(transform_args) do |args|
+          memoized_transforms[args] = Transformation.new(*args)
+        end
+        tree = transformations.fetch(transformation) do |t|
+          transformations[t] = {}
+        end
+        tree[forwarding] = shard
+      end
+
+      transformations
+    end
+  end
+
+  class TransformCommand < BaseTransformCommand
+    def get_transformations
       help!("must have an even number of arguments") unless @argv.length % 2 == 0
 
       scheduler_options = command_options.scheduler_options || {}
       manifest          = manager.manifest(*global_options.tables)
       copy_wrapper      = scheduler_options[:copy_wrapper]
       skip_copies       = scheduler_options[:skip_copies] || false
-      be_quiet          = global_options.force && command_options.quiet
       transformations   = {}
 
-      scheduler_options[:quiet] = be_quiet
-
       @argv.each_slice(2) do |(from_template_s, to_template_s)|
         from, to       = [from_template_s, to_template_s].map {|s| ShardTemplate.parse(s) }
         transformation = Transformation.new(from, to, copy_wrapper, skip_copies)
@@ -808,49 +834,19 @@ def run
         transformations[transformation] = trees
       end
 
-      transformations.reject! {|t,trees| t.noop? or trees.empty? }
-
-      if transformations.empty?
-        puts "Nothing to do!"
-        exit
-      end
-
-      base_name = transformations.values.find {|v| v.is_a?(Hash) && !v.values.empty? }.values.find {|v| !v.nil?}.id.table_prefix.split('_').first
-
-      unless be_quiet
-        transformations.sort.each do |transformation, trees|
-          puts transformation.inspect
-          puts "Applied to #{trees.length} shards"
-          #trees.keys.sort.each {|f| puts "  #{f.inspect}" }
-        end
-        puts ""
-      end
-
-      unless global_options.force
-        print "Continue? (y/n) "; $stdout.flush
-        exit unless $stdin.gets.chomp == "y"
-        puts ""
-      end
-
-      Gizzard.schedule! manager,
-                        base_name,
-                        transformations,
-                        scheduler_options
+      transformations
     end
   end
 
-  class RebalanceCommand < Command
-    def run
+  class RebalanceCommand < BaseTransformCommand
+    def get_transformations
       help!("must have an even number of arguments") unless @argv.length % 2 == 0
 
       scheduler_options = command_options.scheduler_options || {}
       manifest          = manager.manifest(*global_options.tables)
       copy_wrapper      = scheduler_options[:copy_wrapper]
-      be_quiet          = global_options.force && command_options.quiet
       transformations   = {}
 
-      scheduler_options[:quiet] = be_quiet
-
       dest_templates_and_weights = {}
 
       @argv.each_slice(2) do |(weight_s, to_template_s)|
@@ -860,39 +856,12 @@ def run
         dest_templates_and_weights[to] = weight
       end
 
-      transformations = global_options.tables.inject({}) do |all, table|
+      global_options.tables.inject({}) do |all, table|
         trees      = manifest.trees.reject {|(f, s)| f.table_id != table }
         rebalancer = Rebalancer.new(trees, dest_templates_and_weights, copy_wrapper)
 
         all.update(rebalancer.transformations) {|t,a,b| a.merge b }
       end
-
-      if transformations.empty?
-        puts "Nothing to do!"
-        exit
-      end
-
-      base_name = transformations.values.first.values.first.id.table_prefix.split('_').first
-
-      unless be_quiet
-        transformations.each do |transformation, trees|
-          puts transformation.inspect
-          puts "Applied to #{trees.length} shards"
-          #trees.keys.sort.each {|f| puts "  #{f.inspect}" }
-        end
-        puts ""
-      end
-
-      unless global_options.force
-        print "Continue? (y/n) "; $stdout.flush
-        exit unless $stdin.gets.chomp == "y"
-        puts ""
-      end
-
-      Gizzard.schedule! manager,
-                        base_name,
-                        transformations,
-                        scheduler_options
     end
   end
 

diff --git a/lib/gizzmo.rb b/lib/gizzmo.rb
@@ -321,7 +321,7 @@ def add_scheduler_opts(subcommand_options, opts)
     end
   end,
   'transform-tree' => OptionParser.new do |opts|
-    opts.banner = "Usage: #{zero} transform-tree [options] TEMPLATE ROOT_SHARD_ID"
+    opts.banner = "Usage: #{zero} transform-tree [options] TEMPLATE ROOT_SHARD_ID ..."
     separators(opts, DOC_STRINGS['transform-tree'])
 
     add_scheduler_opts subcommand_options, opts