feat: eager include lazy documents during the import

lib/esse/document.rb

@@ -68,9 +68,9 @@ def to_h
     def to_bulk(data: true, operation: nil)
       doc_header.tap do |h|
         if data && operation == :update
-          h[:data] = { doc: source }
+          h[:data] = { doc: source_with_lazy_attributes }
         elsif data
-          h[:data] = source
+          h[:data] = source_with_lazy_attributes
         end
         h.merge!(meta)
       end
@@ -96,5 +96,29 @@ def doc_header
         h[:_routing] = routing if routing?
       end
     end
+
+    def inspect
+      attributes = %i[id routing source].map do |attr|
+        value = send(attr)
+        "#{attr}: #{value.inspect}" if value
+      end.compact.join(', ')
+      "#<#{self.class.name || 'Esse::Document'} #{attributes}>"
+    end
+
+    protected
+
+    def source_with_lazy_attributes
+      return source unless @__lazy_source_data__
+
+      @__lazy_source_data__.merge(source)
+    end
+
+    # api private
+    def __add_lazy_data_to_source__(hash)
+      return hash unless hash.is_a?(Hash)
+
+      @__lazy_source_data__ ||= {}
+      @__lazy_source_data__.merge!(hash)
+    end
   end
 end

lib/esse/index/documents.rb

@@ -214,22 +214,32 @@ def import(*repo_types, context: {}, eager_include_document_attributes: false, l
             cluster.may_update_type!(kwargs)
 
             doc_attrs = {eager: [], lazy: []}
-            # @TODO Implement this
-            if (_expected = eager_include_document_attributes) != false
-              Esse.logger.warn('eager_include_document_attributes is not implemented yet')
-              # allowed = repo.lazy_document_attributes.keys
-              # doc_attrs[:eager] = (expected == true) ? allowed : Array(expected).map(&:to_s) & allowed
+            if (expected = eager_include_document_attributes) != false
+              allowed = repo.lazy_document_attributes.keys
+              doc_attrs[:eager] = (expected == true) ? allowed : Array(expected).map(&:to_s) & allowed
             end
             if (expected = lazy_update_document_attributes) != false
               allowed = repo.lazy_document_attributes.keys
               doc_attrs[:lazy] = (expected == true) ? allowed : Array(expected).map(&:to_s) & allowed
               doc_attrs[:lazy] -= doc_attrs[:eager]
             end
 
+            doc_attrs[:eager].each do |attr_name|
+              partial_docs = repo.documents_for_lazy_attribute(attr_name, *batch.reject(&:ignore_on_index?))
+              next if partial_docs.empty?
+
+              partial_docs.each do |part_doc|
+                doc = batch.find { |d| part_doc.id == d.id && part_doc.type == d.type && part_doc.routing == d.routing }
+                next unless doc
+
+                doc.send(:__add_lazy_data_to_source__, part_doc.source)
+              end
+            end
+
             bulk(**kwargs, index: batch)
 
             doc_attrs[:lazy].each do |attr_name|
-              partial_docs = repo.documents_for_lazy_attribute(attr_name, *batch)
+              partial_docs = repo.documents_for_lazy_attribute(attr_name, *batch.reject(&:ignore_on_index?))
               next if partial_docs.empty?
 
               bulk(**kwargs, update: partial_docs)

spec/support/shared_examples/repository_documents_import.rb

@@ -102,4 +102,44 @@
       end
     end
   end
+
+  context 'when the eager_include_document_attributes is set' do
+    it 'indexes the data and bulk updates all the eager document attributes' do
+      es_client do |client, _conf, cluster|
+        GeosIndex.create_index(alias: true)
+
+        resp = nil
+        expect {
+          resp = GeosIndex::County.import(eager_include_document_attributes: true)
+        }.not_to raise_error
+        expect(resp).to eq(total_counties)
+
+        GeosIndex.refresh
+        expect(GeosIndex.count).to eq(total_counties)
+
+        doc = GeosIndex.get(id: '888')
+        expect(doc.dig('_source', 'country')).to eq('US')
+        expect(doc.dig('_source', 'cities')).to eq([])
+      end
+    end
+
+    it 'indexes the data and bulk updates given eager document attribute' do
+      es_client do |client, _conf, cluster|
+        GeosIndex.create_index(alias: true)
+
+        resp = nil
+        expect {
+          resp = GeosIndex::County.import(eager_include_document_attributes: %i[country])
+        }.not_to raise_error
+        expect(resp).to eq(total_counties)
+
+        GeosIndex.refresh
+        expect(GeosIndex.count).to eq(total_counties)
+
+        doc = GeosIndex.get(id: '888')
+        expect(doc.dig('_source', 'country')).to eq('US')
+        expect(doc.dig('_source', 'cities')).to eq(nil)
+      end
+    end
+  end
 end