add support for `leap facts`. includes some fun new helpers, like run_with_progress(), capture(), and replace_file!().

2 files changed, 125 insertions, 6 deletions

diff --git a/lib/leap_cli/remote/leap_plugin.rb b/lib/leap_cli/remote/leap_plugin.rb
index 2c427e9..8cc96d4 100644
--- a/lib/leap_cli/remote/leap_plugin.rb
+++ b/lib/leap_cli/remote/leap_plugin.rb
@@ -39,6 +39,122 @@ module LeapCli; module Remote; module LeapPlugin
     run "touch #{INITIALIZED_FILE}"
   end
 
+  #
+  # This is a hairy ugly hack, exactly the kind of stuff that makes ruby
+  # dangerous and too much fun for its own good.
+  #
+  # In most places, we run remote ssh without a current 'task'. This works fine,
+  # except that in a few places, the behavior of capistrano ssh is controlled by
+  # the options of the current task.
+  #
+  # We don't want to create an actual current task, because tasks are no fun
+  # and can't take arguments or return values. So, when we need to configure
+  # things that can only be configured in a task, we use this handy hack to
+  # fake the current task.
+  #
+  # This is NOT thread safe, but could be made to be so with some extra work.
+  #
+  def with_task(name)
+    task = @config.tasks[name]
+    @config.class.send(:alias_method, :original_current_task, :current_task)
+    @config.class.send(:define_method, :current_task, Proc.new(){ task })
+    begin
+      yield
+    ensure
+      @config.class.send(:remove_method, :current_task)
+      @config.class.send(:alias_method, :current_task, :original_current_task)
+    end
+  end
+
+  #
+  # similar to run(cmd, &block), but with:
+  #
+  # * exit codes
+  # * stdout and stderr are combined
+  #
+  def stream(cmd, &block)
+    command = '%s 2>&1; echo "exitcode=$?"' % cmd
+    run(command) do |channel, stream, data|
+      exitcode = nil
+      if data =~ /exitcode=(\d+)\n/
+        exitcode = $1.to_i
+        data.sub!(/exitcode=(\d+)\n/,'')
+      end
+      yield({:host => channel[:host], :data => data, :exitcode => exitcode})
+    end
+  end
+
+  #
+  # like stream, but capture all the output before returning
+  #
+  def capture(cmd, &block)
+    command = '%s 2>&1; echo "exitcode=$?" 2>&1;' % cmd
+    host_data = {}
+    run(command) do |channel, stream, data|
+      host_data[channel[:host]] ||= ""
+      if data =~ /exitcode=(\d+)\n/
+        exitcode = $1.to_i
+        data.sub!(/exitcode=(\d+)\n/,'')
+        host_data[channel[:host]] += data
+        yield({:host => channel[:host], :data => host_data[channel[:host]], :exitcode => exitcode})
+      else
+        host_data[channel[:host]] += data
+      end
+    end
+  end
+
+  #
+  # Run a command, with a nice status report and progress indicator.
+  # Only successful results are returned, errors are printed.
+  #
+  # For each successful run on each host, block is yielded with a hash like so:
+  #
+  # {:host => 'bluejay', :exitcode => 0, :data => 'shell output'}
+  #
+  def run_with_progress(cmd, &block)
+    ssh_failures = []
+    exitcode_failures = []
+    succeeded = []
+    task = LeapCli.log_level > 1 ? :standard_task : :skip_errors_task
+    with_task(task) do
+      log :querying, 'facts' do
+        progress "   "
+        call_on_failure do |host|
+          ssh_failures << host
+          progress 'F'
+        end
+        capture(cmd) do |response|
+          if response[:exitcode] == 0
+            progress '.'
+            yield response
+          else
+            exitcode_failures << response
+            progress 'F'
+          end
+        end
+      end
+    end
+    puts "done"
+    if ssh_failures.any?
+      log :failed, 'to connect to nodes: ' + ssh_failures.join(' ')
+    end
+    if exitcode_failures.any?
+      log :failed, 'to run successfully:' do
+        exitcode_failures.each do |response|
+          log "[%s] exit %s - %s" % [response[:host], response[:exitcode], response[:data].strip]
+        end
+      end
+    end
+  rescue Capistrano::RemoteError => err
+    log :error, err.to_s
+  end
+
+  private
+
+  def progress(str='.')
+    $stdout.print str; $stdout.flush;
+  end
+
   #def mkdir(dir)
   #  run "mkdir -p #{dir}"
   #end
diff --git a/lib/leap_cli/remote/tasks.rb b/lib/leap_cli/remote/tasks.rb
index f967db1..0721c34 100644
--- a/lib/leap_cli/remote/tasks.rb
+++ b/lib/leap_cli/remote/tasks.rb
@@ -25,9 +25,12 @@ task :install_prerequisites, :max_hosts => MAX_HOSTS do
   leap.mark_initialized
 end
 
-#task :apply_puppet, :max_hosts => MAX_HOSTS do
-#  raise "now such directory #{puppet_source}" unless File.directory?(puppet_source)
-#  leap.log :applying, "puppet" do
-#    puppet.apply
-#  end
-#end
+#
+# just dummies, used to capture task options
+#
+
+task :skip_errors_task, :on_error => :continue, :max_hosts => MAX_HOSTS do
+end
+
+task :standard_task, :max_hosts => MAX_HOSTS do
+end
\ No newline at end of file