lib/searchkick.rb in searchkick-1.1.1 vs lib/searchkick.rb in searchkick-1.1.2

@@ -6,10 +6,11 @@
 require "searchkick/results"
 require "searchkick/query"
 require "searchkick/reindex_job"
 require "searchkick/model"
 require "searchkick/tasks"
+require "searchkick/middleware"
 require "searchkick/logging" if defined?(Rails)
 
 # background jobs
 begin
   require "active_job"
@@ -25,15 +26,12 @@
   class InvalidQueryError < Elasticsearch::Transport::Transport::Errors::BadRequest; end
   class DangerousOperation < Error; end
   class ImportError < Error; end
 
   class << self
-    attr_accessor :search_method_name
-    attr_accessor :wordnet_path
-    attr_accessor :timeout
-    attr_accessor :models
-    attr_writer :env
+    attr_accessor :search_method_name, :wordnet_path, :timeout, :models
+    attr_writer :client, :env, :search_timeout
   end
   self.search_method_name = :search
   self.wordnet_path = "/var/lib/wn_s.pl"
   self.timeout = 10
   self.models = []
@@ -41,34 +39,95 @@
   def self.client
     @client ||=
       Elasticsearch::Client.new(
         url: ENV["ELASTICSEARCH_URL"],
         transport_options: {request: {timeout: timeout}}
-      )
+      ) do |f|
+        f.use Searchkick::Middleware
+      end
   end
 
-  class << self
-    attr_writer :client
+  def self.env
+    @env ||= ENV["RAILS_ENV"] || ENV["RACK_ENV"] || "development"
   end
 
+  def self.search_timeout
+    @search_timeout || timeout
+  end
+
   def self.server_version
     @server_version ||= client.info["version"]["number"]
   end
 
   def self.enable_callbacks
-    Thread.current[:searchkick_callbacks_enabled] = true
+    self.callbacks_value = nil
   end
 
   def self.disable_callbacks
-    Thread.current[:searchkick_callbacks_enabled] = false
+    self.callbacks_value = false
   end
 
   def self.callbacks?
     Thread.current[:searchkick_callbacks_enabled].nil? || Thread.current[:searchkick_callbacks_enabled]
   end
 
-  def self.env
-    @env ||= ENV["RAILS_ENV"] || ENV["RACK_ENV"] || "development"
+  def self.callbacks(value)
+    if block_given?
+      previous_value = callbacks_value
+      begin
+        self.callbacks_value = value
+        yield
+        perform_bulk if callbacks_value == :bulk
+      ensure
+        self.callbacks_value = previous_value
+      end
+    else
+      self.callbacks_value = value
+    end
+  end
+
+  # private
+  def self.queue_items(items)
+    queued_items.concat(items)
+    perform_bulk unless callbacks_value == :bulk
+  end
+
+  # private
+  def self.perform_bulk
+    items = queued_items
+    clear_queued_items
+    perform_items(items)
+  end
+
+  # private
+  def self.perform_items(items)
+    if items.any?
+      response = client.bulk(body: items)
+      if response["errors"]
+        first_item = response["items"].first
+        raise Searchkick::ImportError, (first_item["index"] || first_item["delete"])["error"]
+      end
+    end
+  end
+
+  # private
+  def self.queued_items
+    Thread.current[:searchkick_queued_items] ||= []
+  end
+
+  # private
+  def self.clear_queued_items
+    Thread.current[:searchkick_queued_items] = []
+  end
+
+  # private
+  def self.callbacks_value
+    Thread.current[:searchkick_callbacks_enabled]
+  end
+
+  # private
+  def self.callbacks_value=(value)
+    Thread.current[:searchkick_callbacks_enabled] = value
   end
 end
 
 # TODO find better ActiveModel hook
 ActiveModel::Callbacks.send(:include, Searchkick::Model)