model_command.rb in ronin-1.1.0.rc1

- old
+ new

@@ -18,58 +18,73 @@
 #
 
 require 'ronin/ui/cli/command'
 require 'ronin/database'
 
-require 'set'
-
 module Ronin
   module UI
     module CLI
       #
-      # A base-command for querying and exporting {Model}s.
+      # A base-command for querying {Model}s.
       #
       class ModelCommand < Command
 
-        class_option :csv, :type => :boolean
-        class_option :xml, :type => :boolean
-        class_option :yaml, :type => :boolean
-        class_option :json, :type => :boolean
+        class_option :database, :type => :string,
+                                :aliases => '-D'
 
         #
         # The model to query.
         #
         # @return [DataMapper::Resource]
         #   The model that will be queried.
         #
-        # @since 1.0.0
+        # @since 1.1.0
         #
-        def self.model
-          @model
+        # @api semipublic
+        #
+        def self.query_model
+          @query_model ||= if self.superclass < ModelCommand
+                       self.superclass.query_model
+                     end
         end
 
         #
         # The query options for the command.
         #
-        # @return [Set]
+        # @return [Array<Symbol>]
         #   The query options and their query method names.
         #
         # @since 1.0.0
         #
+        # @api semipublic
+        #
         def self.query_options
-          @query_options ||= Set[]
+          @query_options ||= []
         end
 
         #
-        # Default method performs the query and prints the found resources.
+        # Iterates over the query options for the command.
         #
-        # @since 1.0.0
+        # @yield [name]
+        #   The given block will be passed the names of the query options.
         #
-        def execute
-          Database.setup
+        # @yieldparam [Symbol] name
+        #   The name of a query option.
+        #
+        # @return [Enumerator]
+        #   If no block is given, an Enumerator object will be returned.
+        #
+        # @since 1.1.0
+        #
+        def self.each_query_option(&block)
+          return enum_for(:each_query_option) unless block
 
-          print_resources(query)
+          self.class.ancestors.each do |ancestor|
+            if ancestor < ModelCommand
+              ancestor.query_options.each(&block)
+            end
+          end
         end
 
         protected
 
         #
@@ -79,14 +94,16 @@
         #   The model class.
         #
         # @return [DataMapper::Resource]
         #   The model that will be queried.
         #
-        # @since 1.0.0
+        # @since 1.1.0
         #
-        def self.model=(model)
-          @model = model
+        # @api semipublic
+        #
+        def self.model(model)
+          @query_model = model
         end
 
         #
         # Defines a query option for the command.
         #
@@ -96,111 +113,79 @@
         # @param [Hash] options
         #   Additional options.
         #
         # @since 1.0.0
         #
+        # @api semipublic
+        #
         def self.query_option(name,options={})
           query_options << name
           class_option(name,options)
         end
 
         #
-        # Invokes a query method on a query.
+        # Sets up the {Database}.
         #
-        # @param [DataMapper::Collection] query
-        #   The query.
+        # @since 1.1.0
         #
-        # @param [Symbol] name
-        #   The method name to call.
+        # @api semipublic
         #
-        # @param [Array] arguments
-        #   Additional arguments to pass to the query method.
-        #
-        # @return [DataMapper::Collection]
-        #   The modified query.
-        #
-        # @since 1.0.0
-        #
-        def query_method(query,name,arguments=[])
-          query_method = begin
-                           query.model.method(name)
-                         rescue NameError
-                           raise("Undefined query method #{query.model}.#{name}")
-                         end
+        def setup
+          super
 
-          case query_method.arity
-          when 0
-            query.send(name)
-          when 1
-            query.send(name,arguments)
-          else
-            query.send(name,*arguments)
+          if self.options[:database]
+            Database.repositories[:default] = options[:database]
           end
+
+          Database.setup
         end
 
         #
         # Builds a new query using the options of the command.
         #
-        # @yield [query]
-        #   If a block is given, it will be passed the new query.
-        #
-        # @yieldparam [DataMapper::Collection] query
-        #   The new query.
-        #
         # @return [DataMapper::Collection]
         #   The new query.
         #
+        # @raise [RuntimeError]
+        #   The query option does not map to a query-method or property
+        #   defined in the Model.
+        #
         # @since 1.0.0
         #
+        # @api semipublic
+        #
         def query
-          new_query = self.class.model.all
-
-          self.class.ancestors.each do |ancestor|
-            if ancestor < ModelCommand
-              ancestor.query_options.each do |name|
-                unless options[name].nil?
-                  new_query = query_method(new_query,name,options[name])
-                end
-              end
-            end
+          unless self.class.query_model
+            raise("query model not defined for #{self.class}")
           end
 
-          new_query = yield(new_query) if block_given?
-          return new_query
-        end
+          query = self.class.query_model.all
 
-        #
-        # Default method which will print every queried resource.
-        #
-        # @param [DataMapper::Resource] resource
-        #   A queried resource from the Database.
-        #
-        # @since 1.0.0
-        #
-        def print_resource(resource)
-          puts resource
-        end
+          self.class.each_query_option do |name|
+            value = options[name]
 
-        #
-        # Prints multiple resources.
-        #
-        # @param [DataMapper::Collection, Array<DataMapper::Resource>] resources
-        #   The query to print.
-        #
-        # @since 1.0.0
-        #
-        def print_resources(resources)
-          if options.csv?
-            print resources.to_csv
-          elsif options.xml?
-            print resources.to_xml
-          elsif options.yaml?
-            print resources.to_yaml
-          elsif options.json?
-            print resources.to_json
-          else
-            resources.each { |resource| print_resource(resource) }
+            # skip unset options
+            next if value.nil?
+
+            if query_model.properties.named?(name)
+              query = query.all(name => value)
+            elsif query_model.respond_to?(name)
+              query_method = query_model.method(name)
+
+              query = case query_method.arity
+                      when 0
+                        query.send(name)
+                      when 1
+                        query.send(name,value)
+                      else
+                        query.send(name,*value)
+                      end
+            else
+              raise("unknown query method or property #{name} for #{query_model}")
+            end
           end
+
+          return query
         end
 
       end
     end
   end