caoutsearch 0.0.1 → 0.0.4

data/lib/caoutsearch/instrumentation/search.rb

@@ -8,10 +8,19 @@ module Caoutsearch
         log_response("Search", event)
       end
 
+      def search_after(event)
+        log_request("SearchAfter", event, format: "truncated")
+        log_response("SearchAfter", event) do |message|
+          payload = event.payload
+          message += ", progress: #{payload[:progress]} / #{payload[:total]}"
+          message
+        end
+      end
+
       def scroll_search(event)
         log_request("Search", event, format: log_request_format)
         log_response("Search", event) do |message|
-          payload  = event.payload
+          payload = event.payload
           message += ", progress: #{payload[:progress]} / #{payload[:total]}"
           message
         end
@@ -20,7 +29,7 @@ module Caoutsearch
       def scroll(event)
         log_request("Scroll", event, format: "truncated")
         log_response("Scroll", event) do |message|
-          payload  = event.payload
+          payload = event.payload
           message += ", progress: #{payload[:progress]} / #{payload[:total]}"
           message
         end

data/lib/caoutsearch/mappings.rb

@@ -18,7 +18,7 @@ module Caoutsearch
     alias_method :as_json, :to_hash
 
     def to_json(*)
-      as_json.to_json
+      MultiJson.dump(as_json)
     end
 
     def include?(*args)

data/lib/caoutsearch/model/indexable.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Model
+    module Indexable
+      extend ActiveSupport::Concern
+
+      included do
+        class_attribute :index_engine_class
+      end
+
+      class_methods do
+        def index_with(index_class)
+          self.index_engine_class = index_class
+
+          extend ClassMethods
+          include InstanceMethods
+        end
+      end
+
+      module ClassMethods
+        def reindex(*properties, **options, &block)
+          index_engine_class.reindex(all, *properties, **options, &block)
+        end
+
+        def delete_indexes
+          find_in_batches do |batch|
+            ids = batch.map(&:id)
+            index_engine_class.delete_documents(ids)
+          end
+        end
+
+        def delete_index(id)
+          index_engine_class.delete_document(id)
+        end
+      end
+
+      module InstanceMethods
+        def as_indexed_json(*properties)
+          index_engine_class.new(self).as_json(*properties)
+        end
+
+        def indexed_document
+          index_engine_class.new(self).indexed_document
+        end
+
+        def update_index(*properties)
+          index_engine_class.new(self).update_document(*properties)
+        end
+
+        def delete_index
+          index_engine_class.new(self).delete_document
+        end
+      end
+    end
+  end
+end

data/lib/caoutsearch/model/searchable.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Model
+    module Searchable
+      extend ActiveSupport::Concern
+
+      included do
+        class_attribute :search_engine_class
+      end
+
+      class_methods do
+        def search_with(search_class)
+          self.search_engine_class = search_class
+
+          extend ClassMethods
+        end
+      end
+
+      module ClassMethods
+        def search_engine
+          search_engine_class.new
+        end
+
+        def search(...)
+          search_engine_class.new.search(...)
+        end
+      end
+    end
+  end
+end

data/lib/caoutsearch/model.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Model
+    extend ActiveSupport::Concern
+
+    included do
+      include Indexable
+      include Searchable
+    end
+  end
+end

data/lib/caoutsearch/response/aggregations.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Response
+    class Aggregations < Caoutsearch::Response::Response
+      disable_warnings
+
+      def initialize(aggs, search = nil, &block)
+        @source_search = search
+        super(aggs, &block)
+      end
+
+      def key?(key)
+        super || has_transformation?(key)
+      end
+
+      def custom_reader(key)
+        if has_transformation?(key)
+          transformation_cache(key) do
+            call_transformation(key)
+          end
+        else
+          super
+        end
+      end
+
+      alias_method :[], :custom_reader
+
+      def has_transformation?(key)
+        @source_search&.class&.transformations&.include?(key.to_s)
+      end
+
+      private
+
+      def call_transformation(key)
+        item = @source_search.class.transformations.fetch(key.to_s)
+        aggs = @source_search.response.aggregations
+
+        @source_search.instance_exec(aggs, &item.block)
+      end
+
+      def transformation_cache(name)
+        @cached_transformations ||= {}
+        return @cached_transformations[name] if @cached_transformations.include?(name)
+
+        @cached_transformations[name] = yield
+      end
+    end
+  end
+end

data/lib/caoutsearch/response/response.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Response
+    class Response < ::Hashie::Mash
+      disable_warnings
+    end
+  end
+end

data/lib/caoutsearch/response/suggestions.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Response
+    class Suggestions < Caoutsearch::Response::Response
+      disable_warnings
+    end
+  end
+end

data/lib/caoutsearch/response.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Response
+  end
+end

data/lib/caoutsearch/search/adapter/active_record.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Search
+    module Adapter
+      module ActiveRecord
+        def self.call(model, hits, skip_query_cache: false)
+          ids = hits.map { |hit| hit["_id"] }
+
+          relation = model.where(model.primary_key => ids).extending(Relation)
+          relation.skip_query_cache! if skip_query_cache
+          relation.hits = hits
+          relation
+        end
+
+        module Relation
+          attr_reader :hits
+
+          def hits=(values)
+            @hits = values
+          end
+
+          # Re-order records based on hits order
+          #
+          def records
+            return super if order_values.present? || @_reordered_records
+
+            load
+            indexes = @hits.each_with_index.to_h { |hit, index| [hit["_id"].to_s, index] }
+            @records = @records.sort_by { |record| indexes[record.id.to_s] }.freeze
+            @_reordered_records = true
+
+            @records
+          end
+        end
+      end
+    end
+  end
+end

data/lib/caoutsearch/search/base.rb

@@ -7,21 +7,22 @@ module Caoutsearch
       include Caoutsearch::Config::Mappings
       include Caoutsearch::Config::Settings
 
-      include Caoutsearch::Search::Search::DeleteMethods
-      include Caoutsearch::Search::Search::Inspect
-      include Caoutsearch::Search::Search::Instrumentation
-      include Caoutsearch::Search::Search::InternalDSL
-      include Caoutsearch::Search::Search::Naming
-      include Caoutsearch::Search::Search::QueryBuilder
-      include Caoutsearch::Search::Search::QueryMethods
-      include Caoutsearch::Search::Search::Resettable
-      include Caoutsearch::Search::Search::Response
-      include Caoutsearch::Search::Search::ScrollMethods
-      include Caoutsearch::Search::Search::SearchMethods
-
-      def self.search(*args)
-        new.search(*args)
-      end
+      include Caoutsearch::Search::BatchMethods
+      include Caoutsearch::Search::Batch::Scroll
+      include Caoutsearch::Search::Batch::SearchAfter
+      include Caoutsearch::Search::Callbacks
+      include Caoutsearch::Search::DeleteMethods
+      include Caoutsearch::Search::Inspect
+      include Caoutsearch::Search::Instrumentation
+      include Caoutsearch::Search::InternalDSL
+      include Caoutsearch::Search::Naming
+      include Caoutsearch::Search::PointInTime
+      include Caoutsearch::Search::QueryBuilder
+      include Caoutsearch::Search::QueryMethods
+      include Caoutsearch::Search::Records
+      include Caoutsearch::Search::Resettable
+      include Caoutsearch::Search::Response
+      include Caoutsearch::Search::SearchMethods
     end
   end
 end

data/lib/caoutsearch/search/batch/scroll.rb

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Search
+    module Batch
+      module Scroll
+        def scroll(scroll: "1m", batch_size: 1000, &block)
+          search = per(batch_size).track_total_hits
+
+          request_payload = {
+            index: index_name,
+            scroll: scroll,
+            body: search.build.to_h
+          }
+
+          total = 0
+          progress = 0
+          requested_at = nil
+          last_response_time = nil
+
+          results = instrument(:scroll_search) do |event_payload|
+            response = client.search(request_payload)
+            last_response_time = Time.current
+
+            total = response["hits"]["total"]["value"]
+            progress += response["hits"]["hits"].size
+
+            event_payload[:request] = request_payload
+            event_payload[:response] = response
+            event_payload[:total] = total
+            event_payload[:progress] = progress
+
+            response
+          end
+
+          hits = results["hits"]["hits"]
+          return if hits.empty?
+
+          yield hits
+          return if progress >= total
+
+          scroll_id = results["_scroll_id"]
+          request_payload = {
+            scroll_id: scroll_id,
+            scroll: scroll
+          }
+
+          loop do
+            requested_at = Time.current
+
+            results = instrument(:scroll, scroll: scroll_id) do |event_payload|
+              response = client.scroll(request_payload)
+              last_response_time = Time.current
+
+              progress += response["hits"]["hits"].size
+
+              event_payload[:request] = request_payload
+              event_payload[:response] = response
+              event_payload[:total] = total
+              event_payload[:progress] = progress
+
+              response
+            rescue Elastic::Transport::Transport::Errors::NotFound => e
+              raise_enhance_message_when_scroll_failed(e, scroll, requested_at, last_response_time)
+            end
+
+            hits = results["hits"]["hits"]
+            break if hits.empty?
+
+            yield hits
+            break if progress >= total
+          end
+
+          total
+        ensure
+          clear_scroll(scroll_id) if scroll_id
+        end
+
+        def clear_scroll(scroll_id)
+          client.clear_scroll(scroll_id: scroll_id)
+        rescue ::Elastic::Transport::Transport::Errors::NotFound
+          # We dont care if the scroll ID is already expired
+        end
+
+        def raise_enhance_message_when_scroll_failed(error, scroll, requested_at, last_response_time)
+          elapsed = (requested_at - last_response_time).round(1).seconds
+
+          raise error.exception("Scroll registered for #{scroll}, #{elapsed.inspect} elapsed between. #{error.message}")
+        end
+      end
+    end
+  end
+end

data/lib/caoutsearch/search/batch/search_after.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Search
+    module Batch
+      module SearchAfter
+        def search_after(keep_alive: "1m", batch_size: 1000, &block)
+          pit_id = open_point_in_time(keep_alive: keep_alive)
+          search = per(batch_size).track_total_hits
+
+          request_payload = {
+            body: search.build.to_h.merge(
+              pit: {
+                id: pit_id,
+                keep_alive: keep_alive
+              }
+            )
+          }
+          request_payload[:body][:sort] ||= [:_shard_doc]
+
+          total = nil
+          progress = 0
+          requested_at = nil
+          last_response_time = Time.current
+
+          loop do
+            requested_at = Time.current
+
+            results = instrument(:search_after, pit: pit_id) do |event_payload|
+              response = client.search(request_payload)
+              last_response_time = Time.current
+
+              total ||= response["hits"]["total"]["value"]
+              progress += response["hits"]["hits"].size
+
+              event_payload[:request] = request_payload
+              event_payload[:response] = response
+              event_payload[:total] = total
+              event_payload[:progress] = progress
+
+              response
+            rescue Elastic::Transport::Transport::Errors::NotFound => e
+              raise_enhance_message_when_pit_failed(e, keep_alive, requested_at, last_response_time)
+            end
+
+            hits = results["hits"]["hits"]
+            pit_id = results["pit_id"]
+            break if hits.empty?
+
+            yield hits
+            break if progress >= total
+
+            request_payload[:body].tap do |body|
+              body[:pit][:id] = pit_id
+              body[:search_after] = hits.last["sort"]
+              body.delete(:track_total_hits)
+            end
+          end
+        ensure
+          close_point_in_time(pit_id) if pit_id
+        end
+
+        def raise_enhance_message_when_pit_failed(error, keep_alive, requested_at, last_response_time)
+          elapsed = (requested_at - last_response_time).round(1).seconds
+
+          raise error.exception("PIT registered for #{keep_alive}, #{elapsed.inspect} elapsed between. #{error.message}")
+        end
+      end
+    end
+  end
+end

data/lib/caoutsearch/search/batch_methods.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Search
+    module BatchMethods
+      def find_each_hit(**options, &block)
+        return to_enum(:find_each_hit, **options) { total_count } unless block
+
+        find_hits_in_batches(**options) do |hits|
+          hits.each(&block)
+        end
+      end
+
+      def find_each_record(**options, &block)
+        return to_enum(:find_each_record, **options) { total_count } unless block
+
+        find_records_in_batches(**options) do |relation|
+          relation.each(&block)
+        end
+      end
+
+      def find_records_in_batches(**options)
+        unless block_given?
+          return to_enum(:find_records_in_batches, **options) do
+            find_hits_in_batches(**options).size
+          end
+        end
+
+        find_hits_in_batches(**options) do |hits|
+          yield records_adapter.call(model, hits, skip_query_cache: true)
+        end
+      end
+
+      def find_hits_in_batches(implementation: :search_after, **options)
+        batch_size = options[:batch_size]&.to_i || @current_limit&.to_i || 1000
+
+        unless block_given?
+          return to_enum(:find_hits_in_batches, **options) do
+            total_count.div(batch_size) + 1
+          end
+        end
+
+        unless %i[search_after scroll].include?(implementation)
+          raise ArgumentError, "unexpected implementation argument: #{implementation.inspect}"
+        end
+
+        method(implementation).call(batch_size: batch_size, **options) do |hits|
+          yield hits
+        end
+      end
+
+      def scroll_records_in_batches(**options)
+        ActiveSupport::Deprecation.warn("scroll_records_in_batches is deprecated, use find_records_in_batches instead")
+        find_records_in_batches(implementation: :scroll, **options)
+      end
+
+      def scroll_records(**options)
+        ActiveSupport::Deprecation.warn("scroll_records is deprecated, use find_each_record instead")
+        find_each_record(implementation: :scroll, **options)
+      end
+    end
+  end
+end

data/lib/caoutsearch/search/callbacks.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Search
+    module Callbacks
+      extend ActiveSupport::Concern
+      include ActiveSupport::Callbacks
+
+      included do
+        define_callbacks :build
+      end
+
+      class_methods do
+        def before_build(*filters, &blk)
+          set_callback(:build, :before, *filters, &blk)
+        end
+
+        def after_build(*filters, &blk)
+          set_callback(:build, :after, *filters, &blk)
+        end
+
+        def around_build(*filters, &blk)
+          set_callback(:build, :around, *filters, &blk)
+        end
+      end
+    end
+  end
+end

data/lib/caoutsearch/search/delete_methods.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Search
+    module DeleteMethods
+      def delete_documents
+        request_payload = {
+          index: index_name,
+          body: build.to_h
+        }
+
+        instrument(:delete) do |event_payload|
+          event_payload[:request] = request_payload
+          event_payload[:response] = client.delete_by_query(request_payload)
+        end
+      end
+    end
+  end
+end

data/lib/caoutsearch/search/dsl/item.rb

@@ -9,9 +9,9 @@ module Caoutsearch
         PROPERTIES_TO_INSPECT = %i[name options block].freeze
 
         def initialize(name, options = {}, &block)
-          @name    = name
+          @name = name
           @options = options
-          @block   = block if block
+          @block = block if block
         end
 
         def has_block?

data/lib/caoutsearch/search/inspect.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Search
+    module Inspect
+      PROPERTIES_TO_INSPECT = %i[
+        search_criteria
+        current_context
+        current_order
+        current_page
+        current_limit
+        current_offset
+        current_aggregations
+        current_suggestions
+        current_returns
+      ].freeze
+
+      def inspect
+        properties = properties_to_inspect.map { |k, v| " #{k}: #{v}" }
+
+        "#<#{self.class}#{properties.join(",")}>"
+      end
+
+      private
+
+      def properties_to_inspect
+        PROPERTIES_TO_INSPECT.each_with_object({}) do |name, properties|
+          value = instance_variable_get("@#{name}")
+          properties[name] = value.inspect if value
+        end
+      end
+    end
+  end
+end

data/lib/caoutsearch/search/instrumentation.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Caoutsearch
+  module Search
+    module Instrumentation
+      extend ActiveSupport::Concern
+
+      def instrument(action, **options, &block)
+        ActiveSupport::Notifications.instrument("#{action}.caoutsearch_search", **options, klass: self.class.to_s, &block)
+      end
+
+      class_methods do
+        def instrument(action, **options, &block)
+          ActiveSupport::Notifications.instrument("#{action}.caoutsearch_search", **options, klass: to_s, &block)
+        end
+      end
+    end
+  end
+end