feat: let the each_serialized_batch eager load attributes

gemfiles/Gemfile.elasticsearch-1.x.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse (0.3.1)
+    esse (0.3.2)
       multi_json
       thor (>= 0.19)
 

gemfiles/Gemfile.elasticsearch-2.x.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse (0.3.1)
+    esse (0.3.2)
       multi_json
       thor (>= 0.19)
 

gemfiles/Gemfile.elasticsearch-5.x.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse (0.3.1)
+    esse (0.3.2)
       multi_json
       thor (>= 0.19)
 

gemfiles/Gemfile.elasticsearch-6.x.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse (0.3.1)
+    esse (0.3.2)
       multi_json
       thor (>= 0.19)
 

gemfiles/Gemfile.elasticsearch-8.x.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse (0.3.1)
+    esse (0.3.2)
       multi_json
       thor (>= 0.19)
 

gemfiles/Gemfile.opensearch-1.x.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse (0.3.1)
+    esse (0.3.2)
       multi_json
       thor (>= 0.19)
 

gemfiles/Gemfile.opensearch-2.x.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse (0.3.1)
+    esse (0.3.2)
       multi_json
       thor (>= 0.19)
 

lib/esse/document.rb

@@ -102,6 +102,7 @@ def inspect
         value = send(attr)
         "#{attr}: #{value.inspect}" if value
       end.compact.join(', ')
+      attributes << " mutations: #{@__mutations__.inspect}" if @__mutations__
       "#<#{self.class.name || 'Esse::Document'} #{attributes}>"
     end
 

lib/esse/index/documents.rb

@@ -209,7 +209,9 @@ def import(*repo_types, context: {}, eager_include_document_attributes: false, l
           doc_attrs[:lazy] = repo.lazy_document_attribute_names(lazy_update_document_attributes)
           doc_attrs[:lazy] -= doc_attrs[:eager]
 
-          repo.each_serialized_batch(**(context || {})) do |batch|
+          context ||= {}
+          context[:attributes] = doc_attrs[:eager] if doc_attrs[:eager].any?
+          repo.each_serialized_batch(**context) do |batch|
             # Elasticsearch 6.x and older have multiple types per index.
             # This gem supports multiple types per index for backward compatibility, but we recommend to update
             # your elasticsearch to a at least 7.x version and use a single type per index.
@@ -219,13 +221,6 @@ def import(*repo_types, context: {}, eager_include_document_attributes: false, l
             kwargs = { suffix: suffix, type: repo_name, **options }
             cluster.may_update_type!(kwargs)
 
-            doc_attrs[:eager].each do |attr_name|
-              repo.retrieve_lazy_attribute_values(attr_name, batch.reject(&:ignore_on_index?)).each do |doc_header, value|
-                doc = batch.find { |d| doc_header.id == d.id && doc_header.type == d.type && doc_header.routing == d.routing }
-                doc&.mutate(attr_name) { value }
-              end
-            end
-
             bulk(**kwargs, index: batch)
 
             doc_attrs[:lazy].each do |attr_name|

lib/esse/repository/lazy_document_attributes.rb

@@ -15,23 +15,23 @@ def lazy_document_attribute_names(all = true)
         when true
           lazy_document_attributes.keys
         else
-          Array(all).map(&:to_s) & lazy_document_attributes.keys
+          filtered = Array(all).map(&:to_s)
+          lazy_document_attributes.keys.select { |name| filtered.include?(name.to_s) }
         end
       end
 
-      def lazy_document_attribute?(attr_name)
-        lazy_document_attributes.key?(attr_name.to_s)
-      end
-
       def fetch_lazy_document_attribute(attr_name)
-        klass, kwargs = lazy_document_attributes.fetch(attr_name.to_s)
+        klass, kwargs = lazy_document_attributes.fetch(attr_name)
         klass.new(**kwargs)
       rescue KeyError
         raise ArgumentError, format('Attribute %<attr>p is not defined as a lazy document attribute', attr: attr_name)
       end
 
       def lazy_document_attribute(attr_name, klass = nil, **kwargs, &block)
-        if lazy_document_attribute?(attr_name)
+        if attr_name.nil?
+          raise ArgumentError, 'Attribute name is required to define a lazy document attribute'
+        end
+        if lazy_document_attribute?(attr_name.to_sym) || lazy_document_attribute?(attr_name.to_s)
           raise ArgumentError, format('Attribute %<attr>p is already defined as a lazy document attribute', attr: attr_name)
         end
 
@@ -40,11 +40,11 @@ def lazy_document_attribute(attr_name, klass = nil, **kwargs, &block)
           klass = Class.new(Esse::DocumentLazyAttribute) do
             define_method(:call, &block)
           end
-          @lazy_document_attributes[attr_name.to_s] = [klass, kwargs]
+          @lazy_document_attributes[attr_name] = [klass, kwargs]
         elsif klass.is_a?(Class) && klass <= Esse::DocumentLazyAttribute
-          @lazy_document_attributes[attr_name.to_s] = [klass, kwargs]
+          @lazy_document_attributes[attr_name] = [klass, kwargs]
         elsif klass.is_a?(Class) && klass.instance_methods.include?(:call)
-          @lazy_document_attributes[attr_name.to_s] = [klass, kwargs]
+          @lazy_document_attributes[attr_name] = [klass, kwargs]
         elsif klass.nil?
           raise ArgumentError, format('A block or a class that responds to `call` is required to define a lazy document attribute')
         else
@@ -53,6 +53,12 @@ def lazy_document_attribute(attr_name, klass = nil, **kwargs, &block)
       ensure
         @lazy_document_attributes&.freeze
       end
+
+      protected
+
+      def lazy_document_attribute?(attr_name)
+        lazy_document_attributes.key?(attr_name)
+      end
     end
 
     extend ClassMethods

lib/esse/repository/object_document_mapper.rb

@@ -74,11 +74,21 @@ def collection(collection_klass = nil, **_, &block)
       # @param [Hash] kwargs The context
       # @return [Enumerator] The enumerator
       # @yield [Array, **context] serialized collection and the optional context from the collection
-      def each_serialized_batch(**kwargs)
+      def each_serialized_batch(attributes: false, **kwargs)
         each_batch(**kwargs) do |*args|
           batch, collection_context = args
           collection_context ||= {}
           entries = [*batch].map { |entry| serialize(entry, **collection_context) }.compact
+          if attributes
+            attributes = lazy_document_attribute_names(attributes) unless attributes.is_a?(Array)
+            attributes.each do |attr_name|
+              retrieve_lazy_attribute_values(attr_name, entries).each do |doc_header, value|
+                doc = entries.find { |d| doc_header.id.to_s == d.id.to_s && doc_header.type == d.type && doc_header.routing == d.routing }
+                doc&.mutate(attr_name) { value }
+              end
+            end
+          end
+
           yield entries, **kwargs
         end
       end

spec/esse/repository/document_spec.rb

@@ -153,6 +153,40 @@
         )
       end
     end
+
+    context 'with lazy_load_attributes' do
+      include_context 'with stories index definition'
+
+      it 'yields serialized objects with lazy attributes when passing attributes: true' do
+        expected_data = []
+        expect {
+          StoriesIndex::Story.each_serialized_batch(attributes: true) do |batch|
+            expected_data.push(*batch)
+          end
+        }.not_to raise_error
+        expect(expected_data.select { |doc| doc.to_h.key?(:tags) && doc.to_h.key?(:tags_count) }).not_to be_empty
+      end
+
+      it 'yields serialized objects without lazy attributes when passing attributes: false' do
+        expected_data = []
+        expect {
+          StoriesIndex::Story.each_serialized_batch(attributes: false) do |batch|
+            expected_data.push(*batch)
+          end
+        }.not_to raise_error
+        expect(expected_data.select { |doc| doc.to_h.key?(:tags) || doc.to_h.key?(:tags_count) }).to be_empty
+      end
+
+      it 'yields serialized objects with lazy attributes when passing specific attributes' do
+        expected_data = []
+        expect {
+          StoriesIndex::Story.each_serialized_batch(attributes: %i[tags]) do |batch|
+            expected_data.push(*batch)
+          end
+        }.not_to raise_error
+        expect(expected_data.select { |doc| doc.to_h.key?(:tags) && !doc.to_h.key?(:tags_count) }).not_to be_empty
+      end
+    end
   end
 
   describe '.documents' do
@@ -391,7 +425,7 @@ def call
     end
 
     it 'returns false' do
-      expect(repo.lazy_document_attribute?(:foo)).to eq(false)
+      expect(repo.send(:lazy_document_attribute?, :foo)).to eq(false)
     end
 
     context 'with a lazy attribute' do
@@ -401,10 +435,10 @@ def call
         end
       end
 
-      it 'returns true for both symbol and string' do
-        expect(repo.lazy_document_attribute?(:foo)).to eq(true)
-        expect(repo.lazy_document_attribute?('foo')).to eq(true)
-        expect(repo.lazy_document_attribute?(:bar)).to eq(false)
+      it 'returns attribute as it is defined' do
+        expect(repo.send(:lazy_document_attribute?, :foo)).to eq(true)
+        expect(repo.send(:lazy_document_attribute?, 'foo')).to eq(false)
+        expect(repo.send(:lazy_document_attribute?, :bar)).to eq(false)
       end
     end
   end
@@ -431,11 +465,11 @@ def call
       end
 
       it 'defines a lazy attribute' do
-        expect(repo.lazy_document_attributes["foo"]).to match_array([
+        expect(repo.lazy_document_attributes[:foo]).to match_array([
           be < Esse::DocumentLazyAttribute,
           {},
         ])
-        expect(repo.lazy_document_attribute?(:foo)).to eq(true)
+        expect(repo.send(:lazy_document_attribute?, :foo)).to eq(true)
       end
     end
 
@@ -470,11 +504,11 @@ def call
       end
 
       it 'defines a lazy attribute' do
-        expect(repo.lazy_document_attributes["foo"]).to match_array([
+        expect(repo.lazy_document_attributes[:foo]).to match_array([
           TheFooParser,
           {},
         ])
-        expect(repo.lazy_document_attribute?(:foo)).to eq(true)
+        expect(repo.send(:lazy_document_attribute?, :foo)).to eq(true)
         expect(repo.lazy_document_attributes).to be_frozen
       end
     end
@@ -487,11 +521,11 @@ def call
       end
 
       it 'defines a lazy attribute' do
-        expect(repo.lazy_document_attributes["foo"]).to match_array([
+        expect(repo.lazy_document_attributes[:foo]).to match_array([
           be < Esse::DocumentLazyAttribute,
           {},
         ])
-        expect(repo.lazy_document_attribute?(:foo)).to eq(true)
+        expect(repo.send(:lazy_document_attribute?, :foo)).to eq(true)
         expect(repo.lazy_document_attributes).to be_frozen
       end
     end
@@ -504,11 +538,11 @@ def call
       end
 
       it 'defines a lazy attribute' do
-        expect(repo.lazy_document_attributes["foo"]).to match_array([
+        expect(repo.lazy_document_attributes[:foo]).to match_array([
           be < Esse::DocumentLazyAttribute,
           { bar: 'baz' },
         ])
-        expect(repo.lazy_document_attribute?(:foo)).to eq(true)
+        expect(repo.send(:lazy_document_attribute?, :foo)).to eq(true)
         expect(repo.lazy_document_attributes).to be_frozen
       end
     end

spec/esse/repository/lazy_document_spec.rb

@@ -12,31 +12,31 @@
     end
 
     it 'returns all the lazy document attribute names as default' do
-      expect(repo.lazy_document_attribute_names).to eq(%w[foo bar])
+      expect(repo.lazy_document_attribute_names).to eq(%i[foo bar])
     end
 
     it 'returns an empty array when no lazy document attributes are defined' do
       expect(Class.new(Esse::Repository).lazy_document_attribute_names).to eq([])
     end
 
     it 'returs all the lazy document attribute names when passing true' do
-      expect(repo.lazy_document_attribute_names(true)).to eq(%w[foo bar])
+      expect(repo.lazy_document_attribute_names(true)).to eq(%i[foo bar])
     end
 
     it 'returns an empty array when passing false' do
       expect(repo.lazy_document_attribute_names(false)).to eq([])
     end
 
     it 'returns an array of lazy document attribute names when passing an array of names' do
-      expect(repo.lazy_document_attribute_names(%w[foo])).to eq(%w[foo])
+      expect(repo.lazy_document_attribute_names(%w[foo])).to eq(%i[foo])
     end
 
     it 'returns an array of lazy document attribute names when passing a single name' do
-      expect(repo.lazy_document_attribute_names('foo')).to eq(%w[foo])
+      expect(repo.lazy_document_attribute_names('foo')).to eq(%i[foo])
     end
 
     it 'returns an array of lazy document attribute names when passing a single name as symbol' do
-      expect(repo.lazy_document_attribute_names(:foo)).to eq(%w[foo])
+      expect(repo.lazy_document_attribute_names(:foo)).to eq(%i[foo])
     end
   end
 

spec/support/shared_contexts/stories_index_definition.rb

@@ -54,6 +54,11 @@
             [doc, ds.find { |s| s[:id] == doc.id.to_i }&.[](:tags) || []]
           end.to_h
         end
+        lazy_document_attribute :tags_count do |docs|
+          docs.map do |doc|
+            [doc, (ds.find { |s| s[:id] == doc.id.to_i }&.[](:tags) || []).size]
+          end.to_h
+        end
       end
     end
   end