lib/sonar/cli/cli.rb in sonar-client-0.1.4 vs lib/sonar/cli/cli.rb in sonar-client-0.1.6

@@ -1,15 +1,14 @@
 # encoding: utf-8
 
 require 'thor'
 require 'sonar/cli/rcfile'
 require 'sonar/search'
+require 'sonar/request'
 require 'awesome_print'
 require 'table_print'
 
-include Sonar::Search
-
 module Sonar
   class CLI < Thor
     class_option 'format', type: :string, desc: 'Flat JSON (include empty collections), JSON lines of collection data (default), or pretty printed [flat/lines/pretty]'
 
     def initialize(*)
@@ -33,14 +32,14 @@
     method_option 'exact', type: :boolean, aliases: '-e', desc: 'Search for the query string exactly, do not include partial string matches'
     def search(type, term)
       types = [type]
 
       if type == 'all'
-        if term =~ IS_IP
-          types = ip_search_type_names
+        if term =~ Search::IS_IP
+          types = @client.ip_search_type_names
         else
-          types = domain_search_type_names
+          types = @client.domain_search_type_names
         end
       end
 
       types.each do |type|
         @query = {}
@@ -53,11 +52,11 @@
     end
 
     desc 'types', 'List all Sonar query types'
     def types
       tp.set :io, $stdout
-      tp QUERY_TYPES, :name, { description: { width: 100 } }, :input
+      tp Search::QUERY_TYPES, :name, { description: { width: 100 } }, :input
     end
 
     desc 'config', 'Sonar config file location'
     def config
       # TODO: add a way to set config
@@ -79,9 +78,24 @@
         end
       else
         # TODO: use a faster JSON generator?
         puts(data.to_json)
       end
+    end
+
+    def handle_search_response(resp)
+      errors = 0
+      if resp.is_a?(Sonar::Request::RequestIterator)
+        resp.each do |data|
+          errors += 1 if data.key?('errors') || data.key?('error')
+          print_json(cleanup_data(data), options['format'])
+        end
+      else
+        errors += 1 if resp.key?('errors') || resp.key?('error')
+        print_json(cleanup_data(resp), options['format'])
+      end
+
+      raise Search::SearchError.new("Encountered #{errors} errors while searching") if errors > 0
     end
 
     # Clean up whitespace and parse JSON values in responses
     def cleanup_data(data)
       return data unless data.is_a?(Hash) && data.has_key?('collection')