Merge pull request #86 from McdonaldSeanp/json_output

mattkirby · web-flow · commit d3d6ebeb1afc · 2020-08-20T14:02:53.000-07:00
Json output for delete/list + better ABS error handling
diff --git a/lib/vmfloaty.rb b/lib/vmfloaty.rb
@@ -87,6 +87,7 @@ def run # rubocop:disable Metrics/AbcSize
       c.option '--verbose', 'Enables verbose output'
       c.option '--service STRING', String, 'Configured pooler service name'
       c.option '--active', 'Prints information about active vms for a given token'
+      c.option '--json', 'Prints information as JSON'
       c.option '--token STRING', String, 'Token for pooler service'
       c.option '--url STRING', String, 'URL of pooler service'
       c.action do |args, options|
@@ -100,10 +101,18 @@ def run # rubocop:disable Metrics/AbcSize
           running_vms = service.list_active(verbose)
           host = URI.parse(service.url).host
           if running_vms.empty?
-            puts "You have no running VMs on #{host}"
+            if options.json
+              puts {}.to_json
+            else
+              FloatyLogger.info "You have no running VMs on #{host}"
+            end
           else
-            puts "Your VMs on #{host}:"
-            Utils.pretty_print_hosts(verbose, service, running_vms)
+            if options.json
+              puts Utils.get_host_data(verbose, service, running_vms).to_json
+            else
+              puts "Your VMs on #{host}:"
+              Utils.pretty_print_hosts(verbose, service, running_vms)
+            end
           end
         else
           # list available vms from pooler
@@ -200,6 +209,7 @@ def run # rubocop:disable Metrics/AbcSize
       c.option '--service STRING', String, 'Configured pooler service name'
       c.option '--all', 'Deletes all vms acquired by a token'
       c.option '-f', 'Does not prompt user when deleting all vms'
+      c.option '--json', 'Outputs hosts scheduled for deletion as JSON'
       c.option '--token STRING', String, 'Token for pooler service'
       c.option '--url STRING', String, 'URL of pooler service'
       c.action do |args, options|
@@ -215,12 +225,18 @@ def run # rubocop:disable Metrics/AbcSize
         if delete_all
           running_vms = service.list_active(verbose)
           if running_vms.empty?
-            puts 'You have no running VMs.'
+            if options.json
+              puts {}.to_json
+            else
+              FloatyLogger.info "You have no running VMs."
+            end
           else
-            Utils.pretty_print_hosts(verbose, service, running_vms, true)
-            # Confirm deletion
             confirmed = true
-            confirmed = agree('Delete all these VMs? [y/N]') unless force
+            unless force
+              Utils.pretty_print_hosts(verbose, service, running_vms, true)
+              # Confirm deletion
+              confirmed = agree('Delete all these VMs? [y/N]')
+            end
             if confirmed
               response = service.delete(verbose, running_vms)
               response.each do |hostname, result|
@@ -257,9 +273,15 @@ def run # rubocop:disable Metrics/AbcSize
 
         unless successes.empty?
           FloatyLogger.info unless failures.empty?
-          puts 'Scheduled the following VMs for deletion:'
-          successes.each do |hostname|
-            puts "- #{hostname}"
+          if options.json
+            puts successes.to_json
+          else
+            puts 'Scheduled the following VMs for deletion:'
+            output = ''
+            successes.each do |hostname|
+              output += "- #{hostname}\n"
+            end
+            puts output
           end
         end
 
diff --git a/lib/vmfloaty/abs.rb b/lib/vmfloaty/abs.rb
@@ -229,10 +229,10 @@ def self.retrieve(verbose, os_types, token, url, user, options, _ondemand = nil)
 
     retries = 360
 
-    raise AuthError, "HTTP #{res.status}: The token provided could not authenticate to the pooler.\n#{res_body}" if res.status == 401
+    validate_queue_status_response(res.status, res.body, "Initial request", verbose)
 
     (1..retries).each do |i|
-      queue_place, res_body = check_queue(conn, saved_job_id, req_obj)
+      queue_place, res_body = check_queue(conn, saved_job_id, req_obj, verbose)
       return translated(res_body) if res_body
 
       sleep_seconds = 10 if i >= 10
@@ -262,11 +262,12 @@ def self.translated(res_body)
     vmpooler_formatted_body
   end
 
-  def self.check_queue(conn, job_id, req_obj)
+  def self.check_queue(conn, job_id, req_obj, verbose)
     queue_info_res = conn.get "status/queue/info/#{job_id}"
     queue_info = JSON.parse(queue_info_res.body)
 
     res = conn.post 'request', req_obj.to_json
+    validate_queue_status_response(res.status, res.body, "Check queue request", verbose)
 
     unless res.body.empty?
       res_body = JSON.parse(res.body)
@@ -315,4 +316,21 @@ def self.disk(_verbose, _url, _hostname, _token, _disk)
   def self.revert(_verbose, _url, _hostname, _token, _snapshot_sha)
     raise NoMethodError, 'revert is not defined for ABS'
   end
+
+  # Validate the http code returned during a queue status request.
+  #
+  # Return a success message that can be displayed if the status code is
+  # success, otherwise raise an error.
+  def self.validate_queue_status_response(status_code, body, request_name, verbose)
+    case status_code
+    when 200
+      "#{request_name} returned success (Code 200)" if verbose
+    when 202
+      "#{request_name} returned accepted, processing (Code 202)" if verbose
+    when 401
+      raise AuthError, "HTTP #{status_code}: The token provided could not authenticate.\n#{body}"
+    else
+      raise "HTTP #{status_code}: #{request_name} request to ABS failed!\n#{body}"
+    end
+  end
 end
diff --git a/lib/vmfloaty/utils.rb b/lib/vmfloaty/utils.rb
@@ -79,40 +79,62 @@ def self.generate_os_hash(os_args)
   end
 
   def self.pretty_print_hosts(verbose, service, hostnames = [], print_to_stderr = false)
+    fetched_data = self.get_host_data(verbose, service, hostnames)
+    fetched_data.each do |hostname, host_data|
+      case service.type
+      when 'ABS'
+      # For ABS, 'hostname' variable is the jobID
+        host_data['allocated_resources'].each do |vm_name, _i|
+          puts "- [JobID:#{host_data['request']['job']['id']}] #{vm_name['hostname']} (#{vm_name['type']}) <#{host_data['state']}>"
+        end
+      when 'Pooler'
+        tag_pairs = []
+        tag_pairs = host_data['tags'].map { |key, value| "#{key}: #{value}" } unless host_data['tags'].nil?
+        duration = "#{host_data['running']}/#{host_data['lifetime']} hours"
+        metadata = [host_data['template'], duration, *tag_pairs]
+        puts "- #{hostname}.#{host_data['domain']} (#{metadata.join(', ')})"
+      when 'NonstandardPooler'
+        line = "- #{host_data['fqdn']} (#{host_data['os_triple']}"
+        line += ", #{host_data['hours_left_on_reservation']}h remaining"
+        line += ", reason: #{host_data['reserved_for_reason']}" unless host_data['reserved_for_reason'].empty?
+        line += ')'
+        puts line
+      else
+        raise "Invalid service type #{service.type}"
+      end
+    end
+  end
+
+  def self.get_host_data(verbose, service, hostnames = [])
+    result = {}
     hostnames = [hostnames] unless hostnames.is_a? Array
     hostnames.each do |hostname|
       begin
         response = service.query(verbose, hostname)
         host_data = response[hostname]
-
-        case service.type
-        when 'ABS'
-          # For ABS, 'hostname' variable is the jobID
-          if host_data['state'] == 'allocated' || host_data['state'] == 'filled'
-            host_data['allocated_resources'].each do |vm_name, _i|
-              puts "- [JobID:#{host_data['request']['job']['id']}] #{vm_name['hostname']} (#{vm_name['type']}) <#{host_data['state']}>"
+        if block_given?
+          yield host_data result
+        else
+          case service.type
+          when 'ABS'
+            # For ABS, 'hostname' variable is the jobID
+            if host_data['state'] == 'allocated' || host_data['state'] == 'filled'
+              result[hostname] = host_data
             end
+          when 'Pooler'
+            result[hostname] = host_data
+          when 'NonstandardPooler'
+            result[hostname] = host_data
+          else
+            raise "Invalid service type #{service.type}"
           end
-        when 'Pooler'
-          tag_pairs = []
-          tag_pairs = host_data['tags'].map { |key, value| "#{key}: #{value}" } unless host_data['tags'].nil?
-          duration = "#{host_data['running']}/#{host_data['lifetime']} hours"
-          metadata = [host_data['template'], duration, *tag_pairs]
-          puts "- #{hostname}.#{host_data['domain']} (#{metadata.join(', ')})"
-        when 'NonstandardPooler'
-          line = "- #{host_data['fqdn']} (#{host_data['os_triple']}"
-          line += ", #{host_data['hours_left_on_reservation']}h remaining"
-          line += ", reason: #{host_data['reserved_for_reason']}" unless host_data['reserved_for_reason'].empty?
-          line += ')'
-          puts line
-        else
-          raise "Invalid service type #{service.type}"
         end
       rescue StandardError => e
         FloatyLogger.error("Something went wrong while trying to gather information on #{hostname}:")
         FloatyLogger.error(e)
       end
     end
+    result
   end
 
   def self.pretty_print_status(verbose, service)
