Result decorator

lib/tire/results/collection.rb

@@ -1,4 +1,37 @@
 module Tire
+  class Exception < ::StandardError; end
+  class UnknownModel < Tire::Exception
+    def initialize(type)
+      @type = type
+    end
+
+    def message
+      "You have tried to eager load the model instances, but " +
+      "Tire cannot find the model class '#{@type.camelize}' " +
+      "based on _type '#{@type}'."
+    end
+  end
+
+  class UnknownType < Tire::Exception
+    def message
+      "You have tried to eager load the model instances, " +
+      "but Tire cannot find the model class because " +
+      "document has no _type property."
+    end
+  end
+
+  class RecordNotFound < Tire::Exception
+    attr_reader :klass, :ids
+
+    def initialize(klass, ids)
+      @klass, @ids = klass, ids
+    end
+
+    def message
+      "Couldn't find all #{klass.name.pluralize} with IDs (#{ids.join(', ')})."
+    end
+  end
+
   module Results
 
     class Collection
@@ -17,52 +50,7 @@ def initialize(response, options={})
       end
 
       def results
-        @results ||= begin
-          hits = @response['hits']['hits'].map { |d| d.update '_type' => Utils.unescape(d['_type']) }
-
-          unless @options[:load]
-            if @wrapper == Hash
-              hits
-            else
-              hits.map do |h|
-                 document = {}
-
-                 # Update the document with content and ID
-                 document = h['_source'] ? document.update( h['_source'] || {} ) : document.update( __parse_fields__(h['fields']) )
-                 document.update( {'id' => h['_id']} )
-
-                 # Update the document with meta information
-                 ['_score', '_type', '_index', '_version', 'sort', 'highlight', '_explanation'].each { |key| document.update( {key => h[key]} || {} ) }
-
-                 # Return an instance of the "wrapper" class
-                 @wrapper.new(document)
-              end
-            end
-
-          else
-            return [] if hits.empty?
-
-            records = {}
-            @response['hits']['hits'].group_by { |item| item['_type'] }.each do |type, items|
-              raise NoMethodError, "You have tried to eager load the model instances, " +
-                                   "but Tire cannot find the model class because " +
-                                   "document has no _type property." unless type
-
-              begin
-                klass = type.camelize.constantize
-              rescue NameError => e
-                raise NameError, "You have tried to eager load the model instances, but " +
-                                 "Tire cannot find the model class '#{type.camelize}' " +
-                                 "based on _type '#{type}'.", e.backtrace
-              end
-              ids = items.map { |h| h['_id'] }
-              records[type] = @options[:load] === true ? klass.find(ids) : klass.find(ids, @options[:load])
-            end
-
-            # Reorder records to preserve order from search results
-            @response['hits']['hits'].map { |item| records[item['_type']].detect { |record| record.id.to_s == item['_id'].to_s } }
-          end
-        end
+        @results ||= fetch_results
       end
 
       def each(&block)
@@ -86,6 +74,7 @@ def to_ary
         self
       end
 
+      private
       # Handles _source prefixed fields properly: strips the prefix and converts fields to nested Hashes
       #
       def __parse_fields__(fields={})
@@ -108,6 +97,67 @@ def __parse_fields__(fields={})
         fields
       end
 
+      def hits
+        @hits ||= @response['hits']['hits'].map { |d| d.update '_type' => Utils.unescape(d['_type']) }
+      end
+
+      def load_records(type, items, options)
+        records = {}
+        if !options.nil?
+          klass = get_class(type)
+          ids = items.map { |h| h['_id'] }
+          (options === true ? klass.find(ids) : klass.find(ids, options)).each do |item|
+            records["#{type}-#{item.id}"] = item
+          end
+        end
+
+        records
+      rescue ActiveRecord::RecordNotFound
+        raise Tire::RecordNotFound.new(klass, ids)
+      end
+
+      def parse_results(type, items)
+        records = load_records(type, items, @options[:load])
+        items.map do |h|
+          document = {}
+
+          # Update the document with content and ID
+          document = h['_source'] ? document.update( h['_source'] || {} ) : document.update( __parse_fields__(h['fields']) )
+          document.update( {'id' => h['_id']} )
+
+          # Update the document with meta information
+          ['_score', '_type', '_index', '_version', 'sort', 'highlight', '_explanation'].each { |key| document.update( {key => h[key]} || {} ) }
+
+          document.update( {'_type' => Utils.unescape(document['_type'])} )
+
+          document['_model'] = records["#{type}-#{h['_id']}"] if @options[:load]
+
+          # Return an instance of the "wrapper" class
+          @wrapper.new(document)
+        end
+      end
+
+      def fetch_results
+        return hits if @wrapper == Hash
+        records = {}
+        hits.group_by { |item| item['_type'] }.each do |type, items|
+          records[type] = parse_results(type, items)
+        end
+
+        sort(records)
+      end
+
+      def get_class(type)
+        raise Tire::UnknownType if type.nil? || type.strip.empty?
+        klass = type.camelize.constantize
+      rescue NameError
+        raise Tire::UnknownModel.new(type)
+      end
+
+      def sort(records)
+        hits.map { |item| records[item['_type']].detect { |record| record.id.to_s == item['_id'].to_s } }
+      end
+
     end
 
   end

lib/tire/results/item.rb

@@ -24,7 +24,11 @@ def initialize(args={})
       # otherwise return +nil+.
       #
       def method_missing(method_name, *arguments)
-        @attributes.has_key?(method_name.to_sym) ? @attributes[method_name.to_sym] : nil
+        if @attributes.has_key?(method_name.to_sym)
+          @attributes[method_name.to_sym]
+        elsif !model.nil?
+          model.send(method_name, *arguments)
+        end
       end
 
       def [](key)
@@ -39,6 +43,10 @@ def type
         @attributes[:_type] || @attributes[:type]
       end
 
+      def model
+        @attributes[:_model]
+      end
+
       def persisted?
         !!id
       end

test/integration/active_record_searchable_test.rb

@@ -95,21 +95,21 @@ def setup
           results = ActiveRecordArticle.search '"Test 1"', :load => true
 
           assert       results.any?
-          assert_equal ActiveRecordArticle.find(1), results.first
+          assert_equal ActiveRecordArticle.find(1), results.first.model
         end
 
         should "load records on block search" do
           results = ActiveRecordArticle.search :load => true do
             query { string '"Test 1"' }
           end
 
-          assert_equal ActiveRecordArticle.find(1), results.first
+          assert_equal ActiveRecordArticle.find(1), results.first.model
         end
 
         should "load records with options on query search" do
           assert_equal ActiveRecordArticle.find(['1'], :include => 'comments').first,
                        ActiveRecordArticle.search('"Test 1"',
-                                                  :load => { :include => 'comments' }).results.first
+                                                  :load => { :include => 'comments' }).results.first.model
         end
 
         should "return empty collection for nonmatching query" do
@@ -121,8 +121,19 @@ def setup
             assert ! results.any?
           end
         end
+
+        should "wrap loaded document with returned record" do
+          results = ActiveRecordArticle.search :load => true do
+            query { string '"Test 1"' }
+            fields :title
+          end
+
+          assert_equal "Test 1", results.first.title
+          assert_equal 6, results.first.length
+        end
+
       end
-      
+
       should "remove document from index on destroy" do
         a = ActiveRecordArticle.new :title => 'Test remove...'
         a.save!
@@ -392,8 +403,8 @@ module ::Rails; end
            # puts s.results[0].inspect
 
            assert_equal 2, s.results.length
-           assert_instance_of ActiveRecordModelOne, s.results[0]
-           assert_instance_of ActiveRecordModelTwo, s.results[1]
+           assert_instance_of ActiveRecordModelOne, s.results[0].model
+           assert_instance_of ActiveRecordModelTwo, s.results[1].model
          end
 
          should "eagerly load all STI descendant records" do
@@ -403,8 +414,8 @@ module ::Rails; end
            end
 
            assert_equal 2, s.results.length
-           assert_instance_of ActiveRecordVideo,  s.results[0]
-           assert_instance_of ActiveRecordPhoto,  s.results[1]
+           assert_instance_of ActiveRecordVideo,  s.results[0].model
+           assert_instance_of ActiveRecordPhoto,  s.results[1].model
          end
       end
 
@@ -434,8 +445,8 @@ module ::Rails; end
           results = ActiveRecordNamespace::MyModel.search 'test', :load => true
 
           assert             results.any?, "No results returned: #{results.inspect}"
-          assert_instance_of ActiveRecordNamespace::MyModel, results.first
-          assert_equal       ActiveRecordNamespace::MyModel.find(1), results.first
+          assert_instance_of ActiveRecordNamespace::MyModel, results.first.model
+          assert_equal       ActiveRecordNamespace::MyModel.find(1), results.first.model
         end
 
       end

test/unit/results_collection_test.rb

@@ -247,14 +247,14 @@ class ResultsCollectionTest < Test::Unit::TestCase
         end
 
         should "raise error when model class cannot be inferred from _type" do
-          assert_raise(NameError) do
+          assert_raise(Tire::UnknownModel) do
             response = { 'hits' => { 'hits' => [ {'_id' => 1, '_type' => 'hic_sunt_leones'}] } }
             Results::Collection.new(response, :load => true).results
           end
         end
 
         should "raise error when _type is missing" do
-          assert_raise(NoMethodError) do
+          assert_raise(Tire::UnknownType) do
             response = { 'hits' => { 'hits' => [ {'_id' => 1}] } }
             Results::Collection.new(response, :load => true).results
           end

test/unit/results_item_test.rb

@@ -122,9 +122,10 @@ class ::FakeRailsModel
             extend  ActiveModel::Naming
             include ActiveModel::Conversion
             def self.find(id, options); new; end
+            def foo; :bar; end
           end
 
-          @document = Results::Item.new :id => 1, :_type => 'fake_rails_model', :title => 'Test'
+          @document = Results::Item.new :id => 1, :_type => 'fake_rails_model', :title => 'Test', :_model => FakeRailsModel.new, :_model => FakeRailsModel.new
         end
 
         should "be an instance of model, based on _type" do
@@ -146,6 +147,10 @@ def self.find(id, options); new; end
           assert_equal Tire::Results::Item, document.class
         end
 
+        should "delegate methods to model" do
+          assert_equal :bar, @document.foo
+        end
+
       end
 
     end