object_document_mapper.rb in esse-0.4.0.rc1

- old
+ new

@@ -67,31 +67,69 @@
         end
 
         @collection_proc = collection_klass || block
       end
 
+      # Expose the collection class to let external plugins and extensions to access it.
+      # @return [Class, nil] The collection class
+      # IDEA: When collection is defined as a block, it should setup a class with the block content.
+      def collection_class
+        return unless @collection_proc.is_a?(Class)
+
+        @collection_proc
+      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(lazy_attributes: false, **kwargs)
+      def each_serialized_batch(eager_load_lazy_attributes: false, preload_lazy_attributes: false, **kwargs)
+        if kwargs.key?(:lazy_attributes)
+          warn 'The `lazy_attributes` option is deprecated. Use `eager_load_lazy_attributes` instead.'
+          eager_load_lazy_attributes = kwargs.delete(:lazy_attributes)
+        end
+
+        lazy_attrs_to_eager_load = lazy_document_attribute_names(eager_load_lazy_attributes)
+        lazy_attrs_to_search_preload = lazy_document_attribute_names(preload_lazy_attributes)
+        lazy_attrs_to_search_preload -= lazy_attrs_to_eager_load
+
         each_batch(**kwargs) do |*args|
           batch, collection_context = args
           collection_context ||= {}
           entries = [*batch].map { |entry| serialize(entry, **collection_context) }.compact
-          if lazy_attributes
-            attrs = lazy_attributes.is_a?(Array) ? lazy_attributes : lazy_document_attribute_names(lazy_attributes)
-            attrs.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 }
+          lazy_attrs_to_eager_load.each do |attr_name|
+            retrieve_lazy_attribute_values(attr_name, entries).each do |doc_header, value|
+              doc = entries.find { |d| d.eql?(doc_header, match_lazy_doc_header: true) }
+              doc&.mutate(attr_name) { value }
+            end
+          end
+
+          if lazy_attrs_to_search_preload.any?
+            entries.group_by(&:routing).each do |routing, docs|
+              search_request = {
+                query: { ids: { values: docs.map(&:id) } },
+                size: docs.size,
+                _source: lazy_attrs_to_search_preload
+              }
+              search_request[:routing] = routing if routing
+              index.search(**search_request).response.hits.each do |hit|
+                header = [hit['_id'], hit['_routing'], hit['_type']]
+                next if header[0].nil?
+
+                hit.dig('_source')&.each do |attr_name, attr_value|
+                  real_attr_name = lazy_document_attribute_names(attr_name).first
+                  next if real_attr_name.nil?
+
+                  doc = entries.find { |d| Esse.document_match_with_header?(d, *header) }
+                  doc&.mutate(real_attr_name) { attr_value }
+                end
               end
             end
           end
 
-          yield entries, **kwargs
+          yield entries
         end
       end
 
       # Wrap collection data into serialized documents
       #
@@ -99,10 +137,10 @@
       #    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|
+          each_serialized_batch(**kwargs) do |docs|
             docs.each { |document| yielder.yield(document) }
           end
         end
       end
 