lib/kriterion/backend/mongodb.rb in kriterion-0.0.1 vs lib/kriterion/backend/mongodb.rb in kriterion-0.1.0

@@ -21,15 +21,14 @@
       attr_reader :sections_db
       attr_reader :items_db
       attr_reader :resources_db
       attr_reader :events_db
       attr_reader :standard_details_db
-      attr_reader :metrics
 
       def initialize(opts)
+        super(opts)
         logger.info 'Initializing MongoDB backend'
-        @metrics             = opts[:metrics] || Kriterion::Metrics.new
         @hostname            = opts[:hostname]
         @port                = opts[:port]
         @database            = opts[:database]
         @client              = Mongo::Client.new(
           ["#{@hostname}:#{@port}"], database: @database
@@ -41,71 +40,40 @@
         @resources_db        = @client[:resources]
         @events_db           = @client[:events]
         @standard_details_db = @client[:standard_details]
       end
 
-      def get_standard(name, opts = {})
-        standard = nil
-        metrics[:backend_get_standard] += Benchmark.realtime do
-          # Set recursion to false by default
-          opts[:recurse] = opts[:recurse] || false
+      def find(type, query, opts)
+        database_for(type).find(query).map do |result|
+          # Sanitise the data if required
+          params = sanitise_data(type, result)
 
-          standard = sanitise_standard(find_standard(name))
-          return nil if standard.nil?
+          # Create the object and return in an array
+          object = class_for(type).new(params)
 
-          find_children!(standard) if opts[:recurse]
-        end
+          # Find children if required
+          find_children!(object) if opts[:recurse]
 
-        standard
-      end
-
-      def find_sections(query)
-        sections_db.find(
-          query
-        ).map do |section|
-          Kriterion::Section.new(section)
+          object
         end
       end
 
-      def add_standard(standard)
-        insert_into_db(standards_db, standard)
+      def insert(object)
+        insert_into_db(database_for(object), object)
       end
 
-      def add_section(section)
-        insert_into_db(sections_db, section)
-      end
-
-      def add_item(item)
-        insert_into_db(items_db, item)
-      end
-
-      def add_resource(resource)
-        insert_into_db(resources_db, resource)
-      end
-
-      def add_event(event)
-        insert_into_db(events_db, event)
-      end
-
       def add_unchanged_node(resource, certname)
         resources_db.update_one(
           { resource: resource.resource },
           '$addToSet' => {
             unchanged_nodes: certname
           }
         )
       end
 
       def update_compliance!(thing)
-        databases = {
-          Kriterion::Standard => standards_db,
-          Kriterion::Section  => sections_db,
-          Kriterion::Item     => items_db,
-          Kriterion::Resource => resources_db
-        }
-
-        databases[thing.class].update_one(
+        database_for(thing).update_one(
           { uuid: thing.uuid },
           '$set' => {
             compliance: thing.compliance
           }
         )
@@ -126,10 +94,46 @@
         )
       end
 
       private
 
+      def sanitise_data(type, data)
+        if type == :standard
+          return nil if data.nil?
+          # Compile the regex from a lazy-compiled BSON regex back to a ruby one
+          data['item_syntax'] = data['item_syntax'].compile
+        end
+        data  
+      end
+
+      # Returns the database for a given object type
+      def database_for(object)
+        cls = class_for(object)
+        databases = {
+          Kriterion::Standard => @standards_db,
+          Kriterion::Section  => @sections_db,
+          Kriterion::Item     => @items_db,
+          Kriterion::Resource => @resources_db,
+          Kriterion::Event    => @events_db
+        }
+        databases[cls]
+      end
+
+      def class_for(name)
+        classes = {
+          'standard' => Kriterion::Standard,
+          'section'  => Kriterion::Section,
+          'item'     => Kriterion::Item,
+          'resource' => Kriterion::Resource,
+          'event'    => Kriterion::Event
+        }
+        # If someone has passed in an object, just return the class
+        return name.class if classes.value? name.class
+
+        classes[name.to_s]
+      end
+
       def find_children!(object)
         accepted_objects = [
           Kriterion::Standard,
           Kriterion::Section,
           Kriterion::Item,
@@ -206,30 +210,9 @@
           parent_uuid: parent.uuid
         )
         result.map do |section|
           Kriterion::Section.new(section)
         end
-      end
-
-      def find_standard(name)
-        result = standards_db.find(name: name)
-        count  = result.count
-        case count
-        when 0
-          nil
-        when 1
-          result.first
-        else
-          raise "Found > 1 standards with name: #{name}"
-        end
-      end
-
-      # Takes a result and sanities it to Kriterion::Standard object
-      def sanitise_standard(result)
-        return nil if result.nil?
-        # Compile the regex from a lazy-compiled BSON regex back to a ruby one
-        result['item_syntax'] = result['item_syntax'].compile
-        Kriterion::Standard.new(result)
       end
     end
   end
 end