WIP: Add federated search

Author: ellnix

lib/meilisearch/rails/multi_search.rb

@@ -1,4 +1,5 @@
-require_relative 'multi_search/result'
+require_relative 'multi_search/multi_search_result'
+require_relative 'multi_search/federated_search_result'
 
 module Meilisearch
   module Rails
@@ -9,23 +10,58 @@ def multi_search(searches)
           normalize(options, index_target)
         end
 
-        MultiSearchResult.new(searches, client.multi_search(search_parameters))
+        MultiSearchResult.new(searches, client.multi_search(queries: search_parameters))
+      end
+
+      def federated_search(queries:, federation: {})
+        if federation.nil?
+          Meilisearch::Rails.logger.warn(
+            '[meilisearch-rails] In federated_search, `nil` is an invalid `:federation` option. To explicitly use defaults, pass `{}`.'
+          )
+
+          federation = {}
+        end
+
+        queries.map! { |item| [nil, item] } if queries.is_a?(Array)
+
+        cleaned_queries = queries.filter_map do |(index_target, options)|
+          index_target = options.delete(:index_uid) || index_target
+          index_target ||= options[:class_name].constantize if options[:class_name]
+
+          strip_pagination_options(options)
+          normalize(options, index_target)
+        end
+
+        FederatedSearchResult.new(queries, client.multi_search(queries: cleaned_queries, federation: federation))
       end
 
       private
 
       def normalize(options, index_target)
+        index_target = index_uid_from_target(index_target)
+
+        return nil if index_target.nil?
+
         options
           .except(:class_name)
-          .merge!(index_uid: index_uid_from_target(index_target))
+          .merge!(index_uid: index_target)
       end
 
       def index_uid_from_target(index_target)
         case index_target
         when String, Symbol
           index_target
-        else
-          index_target.index.uid
+        when Class
+          if index_target.respond_to?(:index)
+            index_target.index.uid
+          else
+            Meilisearch::Rails.logger.warn <<~MODEL_NOT_INDEXED
+              [meilisearch-rails] This class was passed to a multi/federated search but it does not have an #index: #{index_target}
+              [meilisearch-rails] Are you sure it has a `meilisearch` block?
+            MODEL_NOT_INDEXED
+
+            nil
+          end
         end
       end
 
@@ -41,6 +77,20 @@ def paginate(options)
         options[:page] ||= 1
       end
 
+      def strip_pagination_options(options)
+        pagination_options = %w[page hitsPerPage hits_per_page limit].select do |key|
+          options.delete(key) || options.delete(key.to_sym)
+        end
+
+        return if pagination_options.empty?
+
+        Meilisearch::Rails.logger.warn <<~WRONG_PAGINATION
+          [meilisearch-rails] Pagination options in federated search must apply to whole federation.
+          [meilisearch-rails] These options have been removed: #{pagination_options.join(', ')}.
+          [meilisearch-rails] Please pass them after queries, in the `federation:` option.
+        WRONG_PAGINATION
+      end
+
       def pagination_enabled?
         Meilisearch::Rails.configuration[:pagination_backend]
       end

lib/meilisearch/rails/multi_search/federated_search_result.rb

@@ -0,0 +1,76 @@
+require 'active_support/core_ext/module/delegation'
+
+module Meilisearch
+  module Rails
+    class FederatedSearchResult
+      attr_reader :metadata, :hits
+
+      def initialize(searches, raw_results)
+        hits = raw_results.delete('hits')
+        @hits = load_hits(hits, searches.to_a)
+        @metadata = raw_results
+      end
+
+      include Enumerable
+
+      delegate :each, :to_a, :to_ary, to: :@hits
+
+      private
+
+      def load_hits(hits, searches)
+        hits_by_pos = hits.group_by { |hit| hit['_federation']['queriesPosition'] }
+
+        keys_and_records_by_pos = hits_by_pos.to_h do |pos, group_hits|
+          search_target, search_opts = searches[pos]
+
+          klass = if search_opts[:class_name]
+                    search_opts[:class_name].constantize
+                  elsif search_target.instance_of?(Class)
+                    search_target
+                  end
+
+          if klass.present?
+            [pos, load_results(klass, group_hits)]
+          else
+            [pos, [nil, group_hits]]
+          end
+        end
+
+        hits.filter_map do |hit|
+          hit_cond_key, recs_by_id = keys_and_records_by_pos[hit['_federation']['queriesPosition']]
+
+          if hit_cond_key.present?
+            record = recs_by_id[hit[hit_cond_key.to_s].to_s]
+            record&.formatted = hit['_formatted']
+            record
+          else
+            hit
+          end
+        end
+      end
+
+      def load_results(klass, hits)
+        pk_method = klass.ms_primary_key_method
+        pk_method = pk_method.in if Utilities.mongo_model?(klass)
+
+        condition_key = pk_is_virtual?(klass, pk_method) ? klass.primary_key : pk_method
+
+        hits_by_id = hits.index_by { |hit| hit[condition_key.to_s] }
+
+        records = klass.where(condition_key => hits_by_id.keys)
+
+        results_by_id = records.index_by do |record|
+          record.send(condition_key).to_s
+        end
+
+        [condition_key, results_by_id]
+      end
+
+      def pk_is_virtual?(model_class, pk_method)
+        model_class.columns
+                   .map(&(Utilities.sequel_model?(model_class) ? :to_s : :name))
+                   .exclude?(pk_method.to_s)
+      end
+    end
+  end
+end