[MODEL] Implement multi-model search

Related: elastic#10, elastic#30, elastic#50, elastic#129, elastic#346

Author: miguelff

elasticsearch-model/lib/elasticsearch/model.rb

@@ -8,10 +8,13 @@
 
 require 'elasticsearch/model/client'
 
+require 'elasticsearch/model/multimodel'
+
 require 'elasticsearch/model/adapter'
 require 'elasticsearch/model/adapters/default'
 require 'elasticsearch/model/adapters/active_record'
 require 'elasticsearch/model/adapters/mongoid'
+require 'elasticsearch/model/adapters/multiple'
 
 require 'elasticsearch/model/importing'
 require 'elasticsearch/model/indexing'
@@ -64,6 +67,50 @@ module Elasticsearch
   module Model
     METHODS = [:search, :mapping, :mappings, :settings, :index_name, :document_type, :import]
 
+
+    # Keeps a registry of the classes that include `Elasticsearch::Model`
+    #
+    class Registry < Array
+
+      # Add the class of a model to the registry
+      #
+      def self.add(klass)
+        __instance.add(klass)
+      end
+
+      # List all the registered models
+      #
+      # @return [Class]
+      #
+      def self.all
+        __instance.models
+      end
+
+      # Returns the unique instance of the registry
+      #
+      # @api private
+      #
+      def self.__instance
+        @instance ||= new
+      end
+
+      def initialize
+        @models = []
+      end
+
+      # Adds a model to the registry
+      #
+      def add(klass)
+        @models << klass
+      end
+
+      # Gets a copy of the registered models
+      #
+      def models
+        @models.dup
+      end
+    end
+
     # Adds the `Elasticsearch::Model` functionality to the including class.
     #
     # * Creates the `__elasticsearch__` class and instance methods, pointing to the proxy object
@@ -119,6 +166,9 @@ class << self
           include Elasticsearch::Model::Importing::ClassMethods
           include Adapter.from_class(base).importing_mixin
         end
+
+        # Add to the registry if it's a class (and not in intermediate module)
+        Registry.add(base) if base.is_a?(Class)
       end
     end
 
@@ -149,6 +199,28 @@ def client=(client)
         @client = client
       end
 
+      # Search across models which include Elasticsearch::Model
+      #
+      # @param query_or_payload [String,Hash,Object] The search request definition
+      #                                              (string, JSON, Hash, or object responding to `to_hash`)
+      # @param models [Array] The list of Model objects to search
+      # @param options [Hash] Optional parameters to be passed to the Elasticsearch client
+      #
+      # @return [Elasticsearch::Model::Response::Response]
+      #
+      # @example Search across specified models
+      #
+      #     Elasticsearch::Model.search('foo', [Author, Article])
+      #
+      # @example Search across all models
+      #
+      #     Elasticsearch::Model.search('foo')
+      #
+      def search(query_or_payload, models=[], options={})
+        models = Multimodel.new(models)
+        request = Searching::SearchRequest.new(models, query_or_payload, options)
+        Response::Response.new(models, request)
+      end
     end
     extend ClassMethods
 

elasticsearch-model/lib/elasticsearch/model/adapters/multiple.rb

@@ -0,0 +1,125 @@
+module Elasticsearch
+  module Model
+    module Adapter
+
+      # An adapter to be used for deserializing results from multiple models, retrieved through
+      # Elasticsearch::Model.search
+      #
+      # @see Elasticsearch::Model.search
+      #
+      module Multiple
+
+        Adapter.register self, lambda { |klass| klass.is_a? Multimodel }
+
+        module Records
+
+          # Returns an Array, which elements are the model instances represented
+          # by the search results.
+          #
+          # This means that if the models queried are a Mixture of ActiveRecord, Mongoid, or
+          # POROs, the elements contained in this array will also be instances of those models
+          #
+          # Ranking of results across multiple indexes is preserved, and queries made to the different
+          # model's datasources are minimal.
+          #
+          # Internally, it gets the results, as ranked by elasticsearch.
+          # Then results are grouped by _type
+          # Then the model corresponding to each _type is queried to retrieve the records
+          # Finally records are rearranged in the same way results were ranked.
+          #
+          # @return [ElasticSearch::Model]
+          #
+          def records
+            @_records ||= begin
+              result = []
+              by_type = __records_by_type
+              __hits.each do |hit|
+                result << by_type[__type(hit)][hit[:_id]]
+              end
+              result.compact
+            end
+          end
+
+          # Returns the record representation of the results retrieved from Elasticsearch, grouped
+          # by model type
+          #
+          # @example
+          # {Series  =>
+          #   {"1"=> #<Series id: 1, series_name: "The Who S01", created_at: "2015-02-23 17:18:28">},
+          #
+          # "Title =>
+          #   {"1"=> #<Title id: 1, name: "Who Strikes Back", created_at: "2015-02-23 17:18:28">}}
+          #
+          # @api private
+          #
+          def __records_by_type
+            array = __ids_by_type.map do |klass, ids|
+              records = __type_records(klass, ids)
+              ids = records.map(&:id).map(&:to_s)
+              [klass, Hash[ids.zip(records)]]
+            end
+            Hash[array]
+          end
+
+          # Returns the records for a specific type
+          #
+          # @api private
+          #
+          def __type_records(klass, ids)
+            if (adapter = Adapter.adapters[ActiveRecord]) && adapter.call(klass)
+              klass.where(klass.primary_key => ids)
+            elsif (adapter = Adapter.adapters[Mongoid]) && adapter.call(klass)
+              klass.where(:id.in => ids)
+            else
+              klass.find(ids)
+            end
+          end
+
+
+          # @return A Hash containing for each type, the ids to retrieve
+          #
+          # @example {Series =>["1"], Title =>["1", "5"]}
+          #
+          # @api private
+          #
+          def __ids_by_type
+            ids_by_type = {}
+            __hits.each do |hit|
+              type = __type(hit)
+              ids_by_type[type] ||= []
+              ids_by_type[type] << hit[:_id]
+            end
+            ids_by_type
+          end
+
+          # Returns the class of the model associated to a certain hit
+          #
+          # A simple class-level memoization over the `_index` and `_type` properties of the hit is applied.
+          # Hence querying the Model Registry is done the minimal amount of times.
+          #
+          # @see Elasticsearch::Model::Registry
+          #
+          # @return Class
+          #
+          # @api private
+          #
+          def __type(hit)
+            @@__types ||= {}
+            @@__types[[hit[:_index], hit[:_type]].join("::")] ||= begin
+              Registry.all.detect { |model| model.index_name == hit[:_index] && model.document_type == hit[:_type] }
+            end
+          end
+
+
+          # Memoizes and returns the hits from the response
+          #
+          # @api private
+          #
+          def __hits
+            @__hits ||= response.response["hits"]["hits"]
+          end
+        end
+      end
+    end
+  end
+end

elasticsearch-model/lib/elasticsearch/model/multimodel.rb

@@ -0,0 +1,43 @@
+module Elasticsearch
+  module Model
+
+    # Wraps a series of models to be used when querying multiple indexes via Elasticsearch::Model.search
+    #
+    # @see Elasticsearch::Model.search
+    #
+    class Multimodel
+
+      attr_reader :models
+
+      # @param models [Class] The list of models across which the search will be performed.
+      def initialize(*models)
+        @models = models.flatten
+        @models = Model::Registry.all if @models.empty?
+      end
+
+      # Get the list of index names used for retrieving documents when doing a search across multiple models
+      #
+      # @return [String] the list of index names used for retrieving documents
+      #
+      def index_name
+        models.map { |m| m.index_name }
+      end
+
+      # Get the list of document types used for retrieving documents when doing a search across multiple models
+      #
+      # @return [String] the list of document types used for retrieving documents
+      #
+      def document_type
+        models.map { |m| m.document_type }
+      end
+
+      # Get the client common for all models
+      #
+      # @return Elasticsearch::Transport::Client
+      #
+      def client
+        Elasticsearch::Model.client
+      end
+    end
+  end
+end