Code Refactor (#17)

* chore: refactoring

* chore: coerce_to_document is an internal method, make it protected

* chore: each_batch is an internal method, make it protected

* chore: serialize is an internal method, make it protected

* chore: refactoring document class to support source mutations for the lazy attributes

* chore: do not convert to doc on eager loading attribute

* chore: refactoring

* chore: keep repo serialize public

* chore: bump version

Gemfile.lock

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

gemfiles/Gemfile.elasticsearch-1.x.lock

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

gemfiles/Gemfile.elasticsearch-2.x.lock

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

gemfiles/Gemfile.elasticsearch-5.x.lock

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

gemfiles/Gemfile.elasticsearch-6.x.lock

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

gemfiles/Gemfile.elasticsearch-7.x.lock

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

gemfiles/Gemfile.elasticsearch-8.x.lock

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

gemfiles/Gemfile.opensearch-1.x.lock

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

gemfiles/Gemfile.opensearch-2.x.lock

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

lib/esse/document.rb

@@ -56,7 +56,7 @@ def source
 
     # @return [Hash] the document data
     def to_h
-      source.merge(
+      mutated_source.merge(
         _id: id,
       ).tap do |hash|
         hash[:_type] = type if type
@@ -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_with_lazy_attributes }
+          h[:data] = { doc: mutated_source }
         elsif data
-          h[:data] = source_with_lazy_attributes
+          h[:data] = mutated_source
         end
         h.merge!(meta)
       end
@@ -105,20 +105,18 @@ def inspect
       "#<#{self.class.name || 'Esse::Document'} #{attributes}>"
     end
 
-    protected
-
-    def source_with_lazy_attributes
-      return source unless @__lazy_source_data__
-
-      @__lazy_source_data__.merge(source)
+    def mutate(key)
+      @__mutations__ ||= {}
+      @__mutations__[key] = yield
+      instance_variable_set(:@__mutated_source__, nil)
     end
 
-    # api private
-    def __add_lazy_data_to_source__(hash)
-      return hash unless hash.is_a?(Hash)
+    protected
+
+    def mutated_source
+      return source unless @__mutations__
 
-      @__lazy_source_data__ ||= {}
-      @__lazy_source_data__.merge!(hash)
+      @__mutated_source__ ||= source.merge(@__mutations__)
     end
   end
 end

lib/esse/index/documents.rb

@@ -205,15 +205,9 @@ def import(*repo_types, context: {}, eager_include_document_attributes: false, l
 
         repo_hash.slice(*repo_types).each do |repo_name, repo|
           doc_attrs = {eager: [], lazy: []}
-          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] = repo.lazy_document_attribute_names(eager_include_document_attributes)
+          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|
             # Elasticsearch 6.x and older have multiple types per index.
@@ -226,14 +220,9 @@ def import(*repo_types, context: {}, eager_include_document_attributes: false, l
             cluster.may_update_type!(kwargs)
 
             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)
+              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
 

lib/esse/repository/documents.rb

@@ -15,6 +15,12 @@ def update_documents_attribute(name, *ids_or_doc_headers, **kwargs)
       end
 
       def documents_for_lazy_attribute(name, *ids_or_doc_headers)
+        retrieve_lazy_attribute_values(name, *ids_or_doc_headers).map do |doc_header, datum|
+          doc_header.to_doc(name => datum)
+        end
+      end
+
+      def retrieve_lazy_attribute_values(name, *ids_or_doc_headers)
         unless lazy_document_attribute?(name)
           raise ArgumentError, <<~MSG
             The attribute `#{name}` is not defined as a lazy document attribute.
@@ -26,17 +32,16 @@ def documents_for_lazy_attribute(name, *ids_or_doc_headers)
         docs = LazyDocumentHeader.coerce_each(ids_or_doc_headers)
         return [] if docs.empty?
 
-        arr = []
         result = fetch_lazy_document_attribute(name).call(docs)
         return [] unless result.is_a?(Hash)
 
-        result.each do |key, datum|
-          doc = docs.find { |d| d == key || d.id == key }
-          next unless doc
-
-          arr << doc.to_doc(name => datum)
+        result.each_with_object({}) do |(key, value), memo|
+          if key.is_a?(LazyDocumentHeader) && (doc = docs.find { |d| d == key || d.id == key.id })
+            memo[doc] = value
+          elsif (doc = docs.find { |d| d.id == key })
+            memo[doc] = value
+          end
         end
-        arr
       end
     end
 

lib/esse/repository/lazy_document_attributes.rb

@@ -8,6 +8,17 @@ def lazy_document_attributes
         @lazy_document_attributes ||= {}.freeze
       end
 
+      def lazy_document_attribute_names(all = true)
+        case all
+        when false
+          []
+        when true
+          lazy_document_attributes.keys
+        else
+          Array(all).map(&:to_s) & lazy_document_attributes.keys
+        end
+      end
+
       def lazy_document_attribute?(attr_name)
         lazy_document_attributes.key?(attr_name.to_s)
       end

lib/esse/repository/object_document_mapper.rb

@@ -38,31 +38,6 @@ def document(klass = nil, &block)
         end
       end
 
-      def coerce_to_document(value)
-        case value
-        when Esse::Document
-          value
-        when Hash
-          Esse::HashDocument.new(value)
-        when NilClass, FalseClass
-          Esse::NullDocument.new
-        else
-          raise ArgumentError, format('%<arg>p is not a valid document. The document should be a hash or an instance of Esse::Document', arg: value)
-        end
-      end
-
-      # Convert ruby object to json by using the document of the given document type.
-      # @param [Object] model The ruby object
-      # @param [Hash] kwargs The context
-      # @return [Esse::Document] The serialized document
-      def serialize(model, **kwargs)
-        if @document_proc.nil?
-          raise NotImplementedError, format('there is no %<t>p document defined for the %<k>p index', t: repo_name, k: index.to_s)
-        end
-
-        @document_proc.call(model, **kwargs)
-      end
-
       # Used to define the source of data. A block is required. And its
       # content should yield an array of each object that should be serialized.
       # The list of arguments will be passed throught the document method.
@@ -94,6 +69,61 @@ def collection(collection_klass = nil, **_, &block)
         @collection_proc = collection_klass || block
       end
 
+      # Wrap collection data into serialized batches
+      #
+      # @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)
+        each_batch(**kwargs) do |*args|
+          batch, collection_context = args
+          collection_context ||= {}
+          entries = [*batch].map { |entry| serialize(entry, **collection_context) }.compact
+          yield entries, **kwargs
+        end
+      end
+
+      # Wrap collection data into serialized documents
+      #
+      # Example:
+      #    GeosIndex.documents(id: 1).first
+      #
+      # @return [Enumerator] All serialized entries
+      def documents(**kwargs)
+        Enumerator.new do |yielder|
+          each_serialized_batch(**kwargs) do |docs, **_collection_kargs|
+            docs.each { |document| yielder.yield(document) }
+          end
+        end
+      end
+
+      # Convert ruby object to json by using the document of the given document type.
+      # @param [Object] model The ruby object
+      # @param [Hash] kwargs The context
+      # @return [Esse::Document] The serialized document
+      def serialize(model, **kwargs)
+        if @document_proc.nil?
+          raise NotImplementedError, format('there is no %<t>p document defined for the %<k>p index', t: repo_name, k: index.to_s)
+        end
+
+        @document_proc.call(model, **kwargs)
+      end
+
+      protected
+
+      def coerce_to_document(value)
+        case value
+        when Esse::Document
+          value
+        when Hash
+          Esse::HashDocument.new(value)
+        when NilClass, FalseClass
+          Esse::NullDocument.new
+        else
+          raise ArgumentError, format('%<arg>p is not a valid document. The document should be a hash or an instance of Esse::Document', arg: value)
+        end
+      end
+
       # Used to fetch all batch of data defined on the collection model.
       # Arguments can be anything. They will just be passed through the block.
       # Useful when the collection depends on scope or any other conditions
@@ -122,34 +152,6 @@ def each_batch(*args, **kwargs, &block)
       rescue LocalJumpError
         raise(SyntaxError, 'block must be explicitly declared in the collection definition')
       end
-
-      # Wrap collection data into serialized batches
-      #
-      # @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)
-        each_batch(**kwargs) do |*args|
-          batch, collection_context = args
-          collection_context ||= {}
-          entries = [*batch].map { |entry| serialize(entry, **collection_context) }.compact
-          yield entries, **kwargs
-        end
-      end
-
-      # Wrap collection data into serialized documents
-      #
-      # Example:
-      #    GeosIndex.documents(id: 1).first
-      #
-      # @return [Enumerator] All serialized entries
-      def documents(**kwargs)
-        Enumerator.new do |yielder|
-          each_serialized_batch(**kwargs) do |docs, **_collection_kargs|
-            docs.each { |document| yielder.yield(document) }
-          end
-        end
-      end
     end
 
     extend ClassMethods

lib/esse/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Esse
-  VERSION = '0.3.0'
+  VERSION = '0.3.1'
 end

spec/esse/document_spec.rb

@@ -172,4 +172,24 @@ def routing
       end
     end
   end
+
+  describe '#mutate' do
+    let(:document_class) do
+      Class.new(described_class) do
+        def source
+          { foo: 'foo' }
+        end
+      end
+    end
+
+    let(:document) { document_class.new(object, **options) }
+
+    it 'adds the given value to the :mutated_source' do
+      expect {
+        document.mutate(:bar) { 'bar' }
+      }.not_to change(document, :source)
+
+      expect(document.send(:mutated_source)).to eq(foo: 'foo', bar: 'bar')
+    end
+  end
 end