From 557fb947810fa8ddcd2fce56888a73d23aab9e2c Mon Sep 17 00:00:00 2001
From: David Padilla <david@easybroker.com>
Date: Tue, 22 Sep 2015 15:31:33 -0500
Subject: [PATCH] Add the :includes option to ActiveRecord adapter that can be
 used to eager load submodels just like you'd do with regular ActiveRecord.

This helps eliminate the N+1 problem when using
elasticsearch-model searches.

Example:

    class Article < ActiveRecord::Base
      belongs_to :author

      # ...

    end

    Article.search(...).records(includes: [:author])

    # Would be similar to using:

    Article.includes(:author).where(...)
---
 .../lib/elasticsearch/model/adapters/active_record.rb  |  7 +++++++
 .../lib/elasticsearch/model/response.rb                |  4 ++--
 .../lib/elasticsearch/model/response/records.rb        |  3 +++
 .../test/unit/adapter_active_record_test.rb            | 10 ++++++++++
 4 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/elasticsearch-model/lib/elasticsearch/model/adapters/active_record.rb b/elasticsearch-model/lib/elasticsearch/model/adapters/active_record.rb
index b0f61bc8a..2d9bb5378 100644
--- a/elasticsearch-model/lib/elasticsearch/model/adapters/active_record.rb
+++ b/elasticsearch-model/lib/elasticsearch/model/adapters/active_record.rb
@@ -10,10 +10,17 @@ module ActiveRecord
                          lambda { |klass| !!defined?(::ActiveRecord::Base) && klass.respond_to?(:ancestors) && klass.ancestors.include?(::ActiveRecord::Base) }
 
         module Records
+          attr_writer :options
+
+          def options
+            @options ||= {}
+          end
+
           # Returns an `ActiveRecord::Relation` instance
           #
           def records
             sql_records = klass.where(klass.primary_key => ids)
+            sql_records = sql_records.includes(self.options[:includes]) if self.options[:includes]
 
             # Re-order records based on the order from Elasticsearch hits
             # by redefining `to_a`, unless the user has called `order()`
diff --git a/elasticsearch-model/lib/elasticsearch/model/response.rb b/elasticsearch-model/lib/elasticsearch/model/response.rb
index efece7143..a08a36460 100644
--- a/elasticsearch-model/lib/elasticsearch/model/response.rb
+++ b/elasticsearch-model/lib/elasticsearch/model/response.rb
@@ -44,8 +44,8 @@ def results
         #
         # @return [Records]
         #
-        def records
-          @records ||= Records.new(klass, self)
+        def records(options = {})
+          @records ||= Records.new(klass, self, options)
         end
 
         # Returns the "took" time
diff --git a/elasticsearch-model/lib/elasticsearch/model/response/records.rb b/elasticsearch-model/lib/elasticsearch/model/response/records.rb
index cd936559b..4638ca689 100644
--- a/elasticsearch-model/lib/elasticsearch/model/response/records.rb
+++ b/elasticsearch-model/lib/elasticsearch/model/response/records.rb
@@ -12,6 +12,8 @@ class Records
 
         delegate :each, :empty?, :size, :slice, :[], :to_a, :to_ary, to: :records
 
+        attr_accessor :options
+
         include Base
 
         # @see Base#initialize
@@ -25,6 +27,7 @@ def initialize(klass, response, options={})
           metaclass = class << self; self; end
           metaclass.__send__ :include, adapter.records_mixin
 
+          self.options = options
           self
         end
 
diff --git a/elasticsearch-model/test/unit/adapter_active_record_test.rb b/elasticsearch-model/test/unit/adapter_active_record_test.rb
index 19fa238f4..335e3bd10 100644
--- a/elasticsearch-model/test/unit/adapter_active_record_test.rb
+++ b/elasticsearch-model/test/unit/adapter_active_record_test.rb
@@ -51,6 +51,16 @@ def ids
         instance.load
       end
 
+      should "load the records with its submodels when using :includes" do
+        klass    = mock('class', primary_key: :some_key, where: @records)
+        @records.expects(:includes).with([:submodel]).at_least_once
+
+        instance = DummyClassForActiveRecord.new
+        instance.expects(:klass).returns(klass).at_least_once
+        instance.options[:includes] = [:submodel]
+        instance.records
+      end
+
       should "reorder the records based on hits order" do
         @records.instance_variable_set(:@records, @records)