es-elasticity 0.2.11 → 0.3.0

data/lib/elasticity/railtie.rb

@@ -1,23 +1,9 @@
+require "elasticity/log_subscriber"
+
 module Elasticity
   class Railtie < Rails::Railtie
     initializer 'elasticity.initialize_logging' do
-      ActiveSupport::Notifications.subscribe(/\.elasticity$/) do |name, start, finish, id, payload|
-        time = (finish - start)*1000
-
-        if logger = Elasticity.config.logger
-          logger.debug "#{name} #{"%.2f" % time}ms #{MultiJson.dump(payload[:args], pretty: Elasticity.config.pretty_json)}"
-
-          if payload[:backtrace].present?
-            bt = Rails.backtrace_cleaner.clean(payload[:backtrace])
-            logger.debug bt[0,4].join("\n")
-          end
-
-          exception, message = payload[:exception]
-          if exception
-            logger.error "#{name} #{exception}: #{message}"
-          end
-        end
-      end
+      LogSubscriber.attach_to(:elasticity)
     end
   end
 end

data/lib/elasticity/search.rb

@@ -1,76 +1,82 @@
 module Elasticity
-  # Search provides a simple interface for defining a search against an Elasticsearch
-  # index and fetching the results in different ways and mappings.
-  #
-  # Example:
-  #   search = Elasticity::Search.new("people", "person", {...})
-  #   search.documents(Person)
-  class Search
-    attr_reader :index, :document_type, :body
-
-    # Creates a new Search definitions for the given index, document_type and criteria. The
-    # search is not performend until methods are called, each method represents a different
-    # way of fetching and mapping the data.
-    #
-    # The body parameter is a hash following the exact same syntax as Elasticsearch's JSON
-    # query language.
-    def initialize(index, document_type, body)
-      @index          = index
-      @document_type  = document_type.freeze
-      @body           = body.freeze
-    end
+  module Search
+    # Elasticity::Search::Definition is a struct that encapsulates all the data specific to one
+    # ElasticSearch search.
+    class Definition
+      attr_accessor :index_name, :document_type, :body
+
+      def initialize(index_name, document_type, body)
+        @index_name    = index_name
+        @document_type = document_type
+        @body          = body
+      end
+
+      def update(body_changes)
+        self.class.new(@index_name, @document_type, @body.deep_merge(body_changes))
+      end
+
+      def to_search_args
+        { index: @index_name, type: @document_type, body: @body }
+      end
 
-    # Execute the search, fetching only ids from Elasticsearch and then mapping the results
-    # into ActiveRecord models using the provided relation.
-    def active_records(relation)
-      return @active_record if defined?(@active_record)
-      response = @index.search(@document_type, @body.merge(_source: false))
-      @active_record = Result.new(response, ActiveRecordMapper.new(relation))
+      def to_msearch_args
+        { index: @index_name, type: @document_type, search: @body }
+      end
     end
 
-    # Execute the search, fetching all documents from the index and mapping the stored attributes
-    # into instances of the provided class. It will call document_klass.new(attrs), where attrs
-    # are the stored attributes.
-    def documents(document_klass)
-      return @documents if defined?(@documents)
-      response = @index.search(@document_type, @body)
-      @documents = Result.new(response, DocumentMapper.new(document_klass))
+    # Elasticity::Search::Facade provides a simple interface for defining a search and provides
+    # different ways of executing it against Elasticsearch. This is usually the main entry point
+    # for search.
+    class Facade
+      attr_accessor :search_definition
+
+      # Creates a new facade for the given search definition, providing a set of helper methods
+      # to trigger different type of searches and results interpretation.
+      def initialize(client, search_definition)
+        @client            = client
+        @search_definition = search_definition
+      end
+
+      # Performs the search using the default search type and returning an iterator that will yield
+      # hash representations of the documents.
+      def document_hashes
+        LazySearch.new(@client, @search_definition)
+      end
+
+      # Performs the search using the default search type and returning an iterator that will yield
+      # each document, converted to the provided document_klass.
+      def documents(document_klass)
+        LazySearch.new(@client, @search_definition) do |hit|
+          document_klass.from_hit(hit)
+        end
+      end
+
+      # Performs the search using the scan search type and the scoll api to iterate over all the documents
+      # as fast as possible. The sort option will be discarded.
+      #
+      # More info: http://www.elasticsearch.org/guide/en/elasticsearch/guide/current/scan-scroll.html
+      def scan_documents(document_klass, **options)
+        ScanCursor.new(@client, @search_definition, document_klass, **options)
+      end
+
+      # Performs the search only fetching document ids using it to load ActiveRecord objects from the provided
+      # relation. It returns the relation matching the objects found on ElasticSearch.
+      def active_records(relation)
+        ActiveRecordProxy.new(@client, @search_definition, relation)
+      end
     end
 
-    # Result is a collection representing the response from a search against an index. It's what gets
-    # returned by any of the Elasticity::Search methods and it provides a lazily-evaluated and
-    # lazily-mapped – using the provided mapper class.
-    #
-    # Example:
-    #
-    #   response = {"took"=>0, "timed_out"=>false, "_shards"=>{"total"=>5, "successful"=>5, "failed"=>0}, "hits"=>{"total"=>2, "max_score"=>1.0, "hits"=>[
-    #     {"_index"=>"my_index", "_type"=>"my_type", "_id"=>"1", "_score"=>1.0, "_source"=> { "id" => 1, "name" => "Foo" },
-    #     {"_index"=>"my_index", "_type"=>"my_type", "_id"=>"2", "_score"=>1.0, "_source"=> { "id" => 2, "name" => "Bar" },
-    #   ]}}
-    #
-    #   class AttributesMapper
-    #     def map(hits)
-    #       hits.map { |h| h["_source"] }
-    #     end
-    #   end
-    #
-    #   r = Result.new(response, AttributesMapper.new)
-    #   r.total # => 2
-    #   r[0]    # => { "id" => 1, "name" => "Foo" }
-    #
-    class Result
+    class LazySearch
       include Enumerable
 
-      def initialize(response, mapper)
-        @response = response
-        @mapper   = mapper
-      end
+      delegate :each, :size, :length, :[], :+, :-, :&, :|, to: :search_results
 
-      delegate :[], :each, :to_ary, :size, :+, :-, to: :mapping
+      attr_accessor :search_definition
 
-      # The total number of entries as returned by ES
-      def total
-        @response["hits"]["total"]
+      def initialize(client, search_definition, &mapper)
+        @client            = client
+        @search_definition = search_definition
+        @mapper            = mapper
       end
 
       def empty?
@@ -81,46 +87,105 @@ module Elasticity
         empty?
       end
 
+      def total
+        response["hits"]["total"]
+      end
+
       def suggestions
-        @response["suggest"] || {}
+        response["hits"]["suggest"] ||= {}
       end
 
-      def mapping
-        return @mapping if defined?(@mapping)
-        hits = Array(@response["hits"]["hits"])
-        @mapping = @mapper.map(hits)
+      def search_results
+        return @search_results if defined?(@search_results)
+
+        hits = response["hits"]["hits"]
+
+        @search_results = if @mapper.nil?
+          hits
+        else
+          hits.map { |hit| @mapper.(hit) }
+        end
+      end
+
+      private
+
+      def response
+        return @response if defined?(@response)
+        @response = @client.search(@search_definition.to_search_args)
       end
     end
 
-    class DocumentMapper
-      def initialize(document_klass)
-        @document_klass = document_klass
+    class ScanCursor
+      include Enumerable
+
+      def initialize(client, search_definition, document_klass, size: 100, scroll: "1m")
+        @client            = client
+        @search_definition = search_definition
+        @document_klass    = document_klass
+        @size              = size
+        @scroll            = scroll
       end
 
-      def map(hits)
-        hits.map do |hit|
-          attrs = hit["_source"].merge(_id: hit['_id'])
+      def empty?
+        total == 0
+      end
 
-          if hit["highlight"]
-            highlighted_attrs = attrs.dup
-            attrs_set = Set.new
+      def blank?
+        empty?
+      end
 
-            hit["highlight"].each do |name, v|
-              name = name.gsub(/\..*\z/, '')
-              next if attrs_set.include?(name)
-              highlighted_attrs[name] = v
-              attrs_set << name
-            end
+      def total
+        search["hits"]["total"]
+      end
 
-            highlighted = @document_klass.new(highlighted_attrs)
-          end
+      def each_batch
+        enumerator.each do |group|
+          yield(group)
+        end
+      end
 
-          @document_klass.new(attrs.merge(highlighted: highlighted))
+      def each
+        enumerator.each do |group|
+          group.each { |doc| yield(doc) }
         end
       end
+
+      private
+
+      def enumerator
+        Enumerator.new do |y|
+          response = search
+
+          loop do
+            response = @client.scroll(scroll_id: response["_scroll_id"], scroll: @scroll)
+            hits     = response["hits"]["hits"]
+            break if hits.empty?
+
+            y << hits.map { |hit| @document_klass.from_hit(hit) }
+          end
+        end
+      end
+
+      def search
+        return @search if defined?(@search)
+        args    = @search_definition.to_search_args
+        args    = args.merge(search_type: 'scan', size: @size, scroll: @scroll)
+        @search = @client.search(args)
+      end
     end
 
-    class ActiveRecordMapper
+    class ActiveRecordProxy
+      def self.from_hits(relation, hits)
+        ids = hits.map { |hit| hit["_id"] }
+
+        if ids.any?
+          id_col = "#{relation.connection.quote_column_name(relation.table_name)}.#{relation.connection.quote_column_name(relation.klass.primary_key)}"
+          relation.where("#{id_col} IN (?)", ids).order("FIELD(#{id_col},#{ids.join(',')})")
+        else
+          relation.none
+        end
+      end
+
       class Relation < ActiveSupport::ProxyObject
         def initialize(relation)
           @relation = relation
@@ -141,57 +206,59 @@ module Elasticity
         end
       end
 
-      def initialize(relation)
-        @relation = Relation.new(relation)
+      def initialize(client, search_definition, relation)
+        @client            = client
+        @search_definition = search_definition.update(_source: false)
+        @relation          = Relation.new(relation)
       end
 
-      def map(hits)
-        ids = hits.map { |h| h["_id"] }
+      def metadata
+        @metadata ||= { total: response["hits"]["total"], suggestions: response["hits"]["suggest"] || {} }
+      end
 
-        if ids.any?
-          id_col = "#{quote(@relation.table_name)}.#{quote(@relation.klass.primary_key)}"
-          @relation.where(id: ids).order("FIELD(#{id_col},#{ids.join(',')})")
-        else
-          @relation.none
-        end
+      def total
+        metadata[:total]
       end
 
-      private
+      def suggestions
+        metadata[:suggestions]
+      end
 
-      def quote(identifier)
-        @relation.connection.quote_column_name(identifier)
+      def method_missing(name, *args, **options, &block)
+        filtered_relation.public_send(name, *args, **options, &block)
       end
-    end
-  end
 
-  class DocumentSearchProxy < BasicObject
-    def initialize(search, document_klass)
-      @search         = search
-      @document_klass = document_klass
-    end
+      private
 
-    def index
-      @search.index
-    end
+      def response
+        @response ||= @client.search(@search_definition.to_search_args)
+      end
 
-    def document_type
-      @search.document_type
+      def filtered_relation
+        return @filtered_relation if defined?(@filtered_relation)
+        @filtered_relation = ActiveRecordProxy.from_hits(@relation, response["hits"]["hits"])
+      end
     end
 
-    def body
-      @search.body
-    end
+    class DocumentProxy < BasicObject
+      def initialize(search, document_klass)
+        @search         = search
+        @document_klass = document_klass
+      end
 
-    def active_records(relation)
-      @search.active_records(relation)
-    end
+      delegate :search_definition, :active_records, to: :@search
 
-    def documents
-      @search.documents(@document_klass)
-    end
+      def documents
+        @search.documents(@document_klass)
+      end
 
-    def method_missing(method_name, *args, &block)
-      documents.public_send(method_name, *args, &block)
+      def scan_documents(**options)
+        @search.scan_documents(@document_klass, **options)
+      end
+
+      def method_missing(method_name, *args, &block)
+        documents.public_send(method_name, *args, &block)
+      end
     end
   end
 end

data/lib/elasticity/strategies.rb

@@ -0,0 +1,15 @@
+module Elasticity
+  module Strategies
+    class IndexError < StandardError
+      attr_reader :index_base_name
+
+      def initialize(index_base_name, message)
+        @index_name = index_name
+        super("#{index_name}: #{message}")
+      end
+    end
+
+    autoload :SingleIndex, "elasticity/strategies/single_index"
+    autoload :AliasIndex,  "elasticity/strategies/alias_index"
+  end
+end

data/lib/elasticity/strategies/alias_index.rb

@@ -0,0 +1,255 @@
+module Elasticity
+  module Strategies
+    # This strategy keeps two aliases that might be mapped to the same index or different index, allowing
+    # runtime changes by simply atomically updating the aliases. For example, look at the remap method
+    # implementation.
+    class AliasIndex
+      STATUSES = [:missing, :ok]
+
+      def initialize(client, index_base_name)
+        @client       = client
+        @main_alias   = index_base_name
+        @update_alias = "#{index_base_name}_update"
+      end
+
+      # Remap allows zero-downtime/zero-dataloss remap of elasticsearch indexes. Here is the overview
+      # of how it works:
+      #
+      # 1. Creates a new index with the new mapping
+      # 2. Update the aliases so that any write goes to the new index and reads goes to both indexes.
+      # 3. Use scan and scroll to iterate over all the documents in the old index, moving them to the
+      #    new index.
+      # 4. Update the aliases so that all operations goes to the new index.
+      # 5. Deletes the old index.
+      #
+      # It does a little bit more to ensure consistency and to handle race-conditions. For more details
+      # look at the implementation.
+      def remap(index_def)
+        main_indexes   = self.main_indexes
+        update_indexes = self.update_indexes
+
+        if main_indexes.size != 1 || update_indexes.size != 1 || main_indexes != update_indexes
+          raise "Index can't be remapped right now, check if another remapping is already happening"
+        end
+
+        new_index      = create_index(index_def)
+        original_index = main_indexes[0]
+
+        begin
+          # Configure aliases so that search includes the old index and the new index, and writes are made to
+          # the new index.
+          @client.index_update_aliases(body: {
+            actions: [
+              { remove: { index: original_index, alias: @update_alias } },
+              { add:    { index: new_index, alias: @update_alias } },
+              { add:    { index: new_index, alias: @main_alias }},
+            ]
+          })
+
+          @client.index_flush(index: original_index)
+          cursor = @client.search index: original_index, search_type: 'scan', scroll: '1m', _source: false, size: 100
+          loop do
+            cursor = @client.scroll(scroll_id: cursor['_scroll_id'], scroll: '1m')
+            hits   = cursor['hits']['hits']
+            break if hits.empty?
+
+            # Fetch documents based on the ids that existed when the migration started, to make sure we only migrate
+            # documents that haven't been deleted.
+            id_docs = hits.map do |hit|
+              { _index: original_index, _type: hit["_type"], _id: hit["_id"] }
+            end
+
+            docs = @client.mget(body: { docs: id_docs }, refresh: true)["docs"]
+            break if docs.empty?
+
+            # Move only documents that still exists on the old index, into the new index.
+            ops = []
+            docs.each do |doc|
+              ops << { index: { _index: new_index, _type: doc["_type"], _id: doc["_id"], data: doc["_source"] } } if doc["found"]
+            end
+
+            @client.bulk(body: ops)
+
+            # Deal with race conditions by removing from the new index any document that doesn't exist in the old index anymore.
+            ops = []
+            @client.mget(body: { docs: id_docs }, refresh: true)["docs"].each_with_index do |new_doc, idx|
+              if docs[idx]["found"] && !new_doc["found"]
+                ops << { delete: { _index: new_index, _type: new_doc["_type"], _id: new_doc["_id"] } }
+              end
+            end
+
+            @client.bulk(body: ops) unless ops.empty?
+          end
+
+          # Update aliases to only point to the new index.
+          @client.index_update_aliases(body: {
+            actions: [
+              { remove: { index: original_index, alias: @main_alias } },
+            ]
+          })
+          @client.index_delete(index: original_index)
+
+        rescue
+          @client.index_update_aliases(body: {
+            actions: [
+              { add:    { index: original_index, alias: @update_alias } },
+              { remove: { index: new_index, alias: @update_alias } },
+            ]
+          })
+
+          @client.index_flush(index: new_index)
+          cursor = @client.search index: new_index, search_type: 'scan', scroll: '1m', size: 100
+          loop do
+            cursor = @client.scroll(scroll_id: cursor['_scroll_id'], scroll: '1m')
+            hits   = cursor['hits']['hits']
+            break if hits.empty?
+
+            # Move all the documents that exists on the new index back to the old index
+            ops = []
+            hits.each do |doc|
+              ops << { index: { _index: original_index, _type: doc["_type"], _id: doc["_id"], data: doc["_source"] } }
+            end
+
+            @client.bulk(body: ops)
+          end
+
+          @client.index_flush(index: original_index)
+          @client.index_update_aliases(body: {
+            actions: [
+              { remove: { index: new_index, alias: @main_alias } },
+            ]
+          })
+          @client.index_delete(index: new_index)
+
+          raise
+        end
+      end
+
+      def status
+        search_exists = @client.index_exists_alias(name: @main_alias)
+        update_exists = @client.index_exists_alias(name: @update_alias)
+
+        case
+        when search_exists && update_exists
+          :ok
+        when !search_exists && !update_exists
+          :missing
+        else
+          :inconsistent
+        end
+      end
+
+      def missing?
+        status == :missing
+      end
+
+      def main_indexes
+        @client.index_get_aliases(index: "#{@main_alias}-*", name: @main_alias).keys
+      rescue Elasticsearch::Transport::Transport::Errors::NotFound
+        []
+      end
+
+      def update_indexes
+        @client.index_get_aliases(index: "#{@main_alias}-*", name: @update_alias).keys
+      rescue Elasticsearch::Transport::Transport::Errors::NotFound
+        []
+      end
+
+      def create(index_def)
+        if missing?
+          index_name = create_index(index_def)
+          @client.index_update_aliases(body: {
+            actions: [
+              { add: { index: index_name, alias: @main_alias } },
+              { add: { index: index_name, alias: @update_alias } },
+            ]
+          })
+        else
+          raise IndexError.new(@main_alias, "index already exists")
+        end
+      end
+
+      def create_if_undefined(index_def)
+        create(index_def) if missing?
+      end
+
+      def delete
+        @client.index_delete(index: "#{@main_alias}-*")
+      end
+
+      def delete_if_defined
+        delete unless missing?
+      end
+
+      def recreate(index_def)
+        delete_if_defined
+        create(index_def)
+      end
+
+      def index_document(type, id, attributes)
+        res = @client.index(index: @update_alias, type: type, id: id, body: attributes)
+
+        if id = res["_id"]
+          [id, res["created"]]
+        else
+          raise IndexError.new(@update_alias, "failed to index document")
+        end
+      end
+
+      def delete_document(type, id)
+        ops = (main_indexes | update_indexes).map do |index|
+          { delete: { _index: index, _type: type, _id: id } }
+        end
+
+        @client.bulk(body: ops)
+      end
+
+      def get_document(type, id)
+        @client.get(index: @main_alias, type: type, id: id)
+      end
+
+      def search(type, body)
+        Search::Facade.new(@client, Search::Definition.new(@main_alias, type, body))
+      end
+
+      def delete_by_query(type, body)
+        @client.delete_by_query(index: @main_alias, type: type, body: body)
+      end
+
+      def bulk
+        b = Bulk::Alias.new(@client, @update_alias, main_indexes)
+        yield b
+        b.execute
+      end
+
+      def flush
+        @client.index_flush(index: @update_alias)
+      end
+
+      def settings
+        args = { index: @main_alias }
+        settings = @client.index_get_settings(index: @main_alias)
+        settings[@main_alias]["settings"]
+      rescue Elasticsearch::Transport::Transport::Errors::NotFound
+        nil
+      end
+
+      def mappings
+        args = { index: @main_alias }
+        mapping = @client.index_get_mapping(index: @main_alias)
+        mapping[@main_alias]["mappings"]
+      rescue Elasticsearch::Transport::Transport::Errors::NotFound
+        nil
+      end
+
+      private
+
+      def create_index(index_def)
+        ts = Time.now.utc.strftime("%Y-%m-%d_%H:%M:%S.%6N")
+        index_name = "#{@main_alias}-#{ts}"
+        @client.index_create(index: index_name, body: index_def)
+        index_name
+      end
+    end
+  end
+end