elasticsearch_record 1.0.0

data/lib/elasticsearch_record/relation/query_methods.rb

@@ -0,0 +1,276 @@
+module ElasticsearchRecord
+  module Relation
+    module QueryMethods
+
+      # unsupported method
+      def joins(*)
+        raise ActiveRecord::StatementInvalid, 'Unsupported method "joins"'
+      end
+
+      # unsupported method
+      # def includes(*)
+      #   raise ActiveRecord::StatementInvalid, 'Unsupported method "includes"'
+      # end
+
+      # sets or overwrites the query kind (e.g. compound queries -> :bool, :boosting, :constant_score, ...).
+      # Also other query kinds like :intervals, :match, ... are allowed.
+      # As an alternative you can also call the #query(<kind>,{argument}) method.
+      # @param [String, Symbol] value - the kind
+      def kind(value)
+        spawn.kind!(value)
+      end
+
+      # same like +#kind+, but on the same relation (no spawn)
+      def kind!(value)
+        # :nodoc:
+        self.kind_value = value
+        self
+      end
+
+      # sets or overwrites additional arguments for the query (not the :query-node, the whole query).
+      # You can also force a overwrite of previously defined arguments, like +size+ or +from+.
+      # This is useful to force remove of keys.
+      #
+      # @example
+      #   # adds {refresh true} to the query
+      #   configure(:refresh, true)
+      #
+      #   # overwrites or sets {from: 50} but removes the :sort key
+      #   configure({from: 50, sort: nil})
+      # @param [Array] args
+      def configure(*args)
+        spawn.configure!(*args)
+      end
+
+      # same like +#configure!+, but on the same relation (no spawn)
+      def configure!(*args)
+        check_if_method_has_arguments!(__callee__, args)
+
+        if args.length == 1 && args.first.is_a?(Hash)
+          self.configure_value = self.configure_value.merge(args[0])
+        elsif args.length == 2 && args[0] == :__claim__
+          tmp = self.configure_value[:__claim__] || []
+          tmp << args[1]
+          self.configure_value = self.configure_value.merge(:__claim__ => tmp)
+        elsif args.length == 2
+          self.configure_value = self.configure_value.merge(args[0] => args[1])
+        end
+
+        self
+      end
+
+      def aggregate(*args)
+        check_if_method_has_arguments!(__callee__, args)
+        spawn.aggregate!(*args)
+      end
+
+      alias_method :aggs, :aggregate
+
+      def aggregate!(opts, *rest)
+        case opts
+        when Symbol, String
+          self.aggs_clause += build_query_clause(opts, rest)
+        when Hash
+          opts.each do |key, value|
+            self.aggs_clause += build_query_clause(key, value)
+          end
+        else
+          raise ArgumentError, "Unsupported argument type for aggregate: #{opts}"
+        end
+
+        self
+      end
+
+      def query(*args)
+        check_if_method_has_arguments!(__callee__, args)
+        spawn.query!(*args)
+      end
+
+      def query!(kind, opts, *rest)
+        kind!(kind)
+        self.query_clause += build_query_clause(opts.keys[0], opts.values[0], rest)
+        self
+      end
+
+      def filter(*args)
+        check_if_method_has_arguments!(__callee__, args)
+        spawn.filter!(*args)
+      end
+
+      def filter!(opts, *rest)
+        # :nodoc:
+        set_default_kind!
+        self.query_clause += build_query_clause(:filter, opts, rest)
+        self
+      end
+
+      def must_not(*args)
+        check_if_method_has_arguments!(__callee__, args)
+        spawn.must_not!(*args)
+      end
+
+      def must_not!(opts, *rest)
+        # :nodoc:
+        set_default_kind!
+        self.query_clause += build_query_clause(:must_not, opts, rest)
+        self
+      end
+
+      def must(*args)
+        check_if_method_has_arguments!(__callee__, args)
+        spawn.must!(*args)
+      end
+
+      def must!(opts, *rest)
+        # :nodoc:
+        set_default_kind!
+        self.query_clause += build_query_clause(:must, opts, rest)
+        self
+      end
+
+      def should(*args)
+        check_if_method_has_arguments!(__callee__, args)
+        spawn.should!(*args)
+      end
+
+      def should!(opts, *rest)
+        # :nodoc:
+        set_default_kind!
+        self.query_clause += build_query_clause(:should, opts, rest)
+        self
+      end
+
+      # creates a condition on the relation.
+      # There are several possibilities to call this method.
+      #
+      #   # create a simple 'term' condition on the query[:filter] param
+      #     where({name: 'hans'})
+      #     > query[:filter] << { term: { name: 'hans' } }
+      #
+      #   # create a simple 'terms' condition on the query[:filter] param
+      #     where({name: ['hans','peter']})
+      #     > query[:filter] << { terms: { name: ['hans','peter'] } }
+      #
+      #     where(:must_not, term: {name: 'horst'})
+      #     where(:query_string, "(new york OR dublin)", fields: ['name','description'])
+      #
+      #   # nested array
+      #   where([ [:filter, {...}], [:must_not, {...}]])
+      def where(*args)
+        return none if args[0] == :none
+
+        super
+      end
+
+      def where!(opts, *rest)
+        case opts
+          # check the first provided parameter +opts+ and validate, if this is an alias for "must, must_not, should or filter"
+          # if true, we expect the rest[0] to be a hash.
+          # For this correlation we forward this as RAW-data without check & manipulation
+        when Symbol
+          case opts
+          when :filter, :must, :must_not, :should
+            send("#{opts}!", *rest)
+          else
+            raise ArgumentError, "Unsupported argument type for where: #{opts}"
+          end
+        when Array
+          # check if this is a nested array of multiple [<kind>,<data>]
+          if opts[0].is_a?(Array)
+            opts.each { |item|
+              where!(*item)
+            }
+          else
+            where!(*opts, *rest)
+          end
+        when String
+          # fallback to ActiveRecords +#where_clause+
+          # currently NOT supported
+          super(opts, rest)
+        else
+          # hash -> {name: 'hans'}
+          # protects against forwarding params directly to where ...
+          # User.where(params) <- will never work
+          # User.where(params.permit(:user)) <- ok
+          opts = sanitize_forbidden_attributes(opts)
+
+          # resolve possible aliases
+          opts = opts.transform_keys do |key|
+            key = key.to_s
+            klass.attribute_aliases[key] || key
+          end
+
+          # check if we have keys without Elasticsearch fields
+          if (invalid = (opts.keys - klass.searchable_column_names)).present?
+            raise(ActiveRecord::UnknownAttributeReference,
+                  "Unable to build query with unknown searchable attributes: #{invalid.map(&:inspect).join(", ")}. " \
+                "If you want to build a custom query you should use one of those methods: 'filter, must, must_not, should'. " \
+                "#{klass.name}.filter('#{invalid[0]}' => '...')"
+            )
+          end
+
+          # force set default kind
+          set_default_kind!
+
+          # builds predicates from opts (transforms this in a more unreadable way but is required for nested assignment & binds ...)
+          parts = predicate_builder.build_from_hash(opts) do |table_name|
+            lookup_table_klass_from_join_dependencies(table_name)
+          end
+
+          self.where_clause += ::ActiveRecord::Relation::WhereClause.new(parts)
+        end
+
+        self
+      end
+
+      def unscope!(*args)
+        # :nodoc:
+        self.unscope_values += args
+
+        args.each do |scope|
+          case scope
+          when Symbol
+            unless _valid_unscoping_values.include?(scope)
+              raise ArgumentError, "Called unscope() with invalid unscoping argument ':#{scope}'. Valid arguments are :#{_valid_unscoping_values.to_a.join(", :")}."
+            end
+            assert_mutability!
+            @values.delete(scope)
+          when Hash
+            scope.each do |key, target_value|
+              target_query_clause = build_query_clause(key, target_value)
+              self.query_clause   -= target_query_clause
+            end
+          else
+            raise ArgumentError, "Unrecognized scoping: #{args.inspect}. Use .unscope(where: :attribute_name) or .unscope(:order), for example."
+          end
+        end
+
+        self
+      end
+
+      private
+
+      def build_query_clause(kind, data, rest = [])
+        ElasticsearchRecord::Relation::QueryClause.new(kind, Array.wrap(data), rest.extract_options!)
+      end
+
+      # sets the default kind if no kind value was defined.
+      # this is called by all conditional methods (where, not, filter, must, must_not & should)
+      def set_default_kind!
+        self.kind_value ||= :bool
+      end
+
+      # overwrite default method to add additional values for kind, query, aggs, ...
+      def build_arel(*args)
+        arel = super(*args)
+
+        arel.kind(kind_value) if kind_value
+        arel.query(query_clause.ast) unless query_clause.empty?
+        arel.aggs(aggs_clause.ast) unless aggs_clause.empty?
+        arel.configure(configure_value) if configure_value.present?
+
+        arel
+      end
+    end
+  end
+end

data/lib/elasticsearch_record/relation/result_methods.rb

@@ -0,0 +1,222 @@
+module ElasticsearchRecord
+  module Relation
+    module ResultMethods
+      # aggregate pluck provided columns.
+      # returns a hash of values for each provided column
+      #
+      # Person.agg_pluck(:name)
+      # => {"name" => ['David', 'Jeremy', 'Jose']}
+      #
+      # Person.agg_pluck(:id, :name)
+      # => {"id" => ['11', '2', '5'], "name" => ['David', 'Jeremy', 'Jose']}
+      #
+      # @param [Array] column_names
+      # @return [Hash]
+      def agg_pluck(*column_names)
+        scope = self.spawn
+
+        column_names.each do |column_name|
+          scope.aggregate!(column_name, { terms: { field: column_name, size: limit_value || 10 } })
+        end
+
+        scope.aggregations.reduce({}) { |m, (k, v)|
+          m[k.to_s] = v[:buckets].map { |bucket| bucket[:key] }
+          m
+        }
+      end
+
+      # A multi-bucket aggregation that creates composite buckets from different sources.
+      # PLEASE NOTE: The composite aggregation is expensive. Load test your application
+      # before deploying a composite aggregation in production!
+      #
+      # For a single column_name a hash with the distinct key and the +doc_count+ as value is returned.
+      # For multiple column_names a hash with the distinct keys (as hash) and the +doc_count+ as value is returned.
+      #
+      # Person.composite(:name)
+      # => {"David" => 10, "Jeremy" => 1, "Jose" => 24}
+      #
+      # Person.composite(:name, :age)
+      # => {
+      #       {name: "David", age: "16"} => 3,
+      #       {name: "David", age: "18"} => 6,
+      #       {name: "David", age: "20"} => 1,
+      #       {name: "Jeremy", age: "20"} => 1,
+      #       {name: "Jose", age: "6"} => 2,
+      #       ...
+      #    }
+      # @param [Array] column_names
+      # @return [Hash]
+      def composite(*column_names)
+        scope = self.spawn
+        scope.aggregate!(:composite_bucket, { composite: { size: limit_value || 10, sources: column_names.map { |column_name| { column_name => { terms: { field: column_name } } } } } })
+
+        if column_names.size == 1
+          column_name = column_names[0]
+          scope.aggregations[:composite_bucket][:buckets].reduce({}) { |m, bucket| m[bucket[:key][column_name]] = bucket[:doc_count]; m }
+        else
+          scope.aggregations[:composite_bucket][:buckets].reduce({}) { |m, bucket| m[bucket[:key]] = bucket[:doc_count]; m }
+        end
+      end
+
+      # creates and returns a new point in time id.
+      # optionally yields the provided block and closes the pit afterwards.
+      # @param [String] keep_alive (default: '1m')
+      # @return [nil, String] - either returns the pit_id (no block given) or nil
+      def point_in_time(keep_alive: '1m')
+        # resolve a initial PIT id
+        initial_pit_id = klass.connection.api(:core, :open_point_in_time, { index: klass.table_name, keep_alive: keep_alive }, "#{klass} Open Pit").dig('id')
+
+        return initial_pit_id unless block_given?
+
+        # block provided, so yield with id
+        yield initial_pit_id
+
+        # close PIT
+        klass.connection.api(:core, :close_point_in_time, { body: { id: initial_pit_id } }, "#{klass} Close Pit")
+
+        # return nil if everything was ok
+        nil
+      end
+
+      alias_method :pit, :point_in_time
+
+      # executes the current query in a +point_in_time+ scope.
+      # this will provide the possibility to resolve more than the +max_result_window+ (default: 10000) hits.
+      # resolves results (hits->hits) from the search but uses the pit query instead to resolve more than 10000 entries.
+      #
+      # If a block was provided it'll yield the results array per batch size.
+      #
+      # @param [String] keep_alive - how long to keep alive (for each single request) - default: '1m'
+      # @param [Integer] batch_size - how many results per query (default: 1000 - this means at least 10 queries before reaching the +max_result_window+)
+      def pit_results(keep_alive: '1m', batch_size: 1000)
+        # check if a limit or offset values was provided
+        results_limit  = limit_value ? limit_value : Float::INFINITY
+        results_offset = offset_value ? offset_value : 0
+
+        # search_after requires a order - we resolve a order either from provided value or by default ...
+        relation = ordered_relation
+
+        # clear limit & offset
+        relation.offset!(nil).limit!(nil)
+
+        # remove the 'index' from the query arguments (pit doesn't like that)
+        relation.configure!(:__claim__, { index: nil })
+
+        # we store the results in this array
+        results       = []
+        results_total = 0
+
+        # resolve a new pit and auto-close after we finished
+        point_in_time(keep_alive: keep_alive) do |pit_id|
+          current_pit_hash = { pit: { id: pit_id, keep_alive: keep_alive } }
+
+          # resolve new data until we got all we need
+          loop do
+            # change pit settings & limit (spawn is required, since a +resolve+ will make the relation immutable)
+            current_response = relation.spawn.configure!(current_pit_hash).limit!(batch_size).resolve('Pit').response
+
+            # resolve only data from hits->hits[{_source}]
+            current_results        = current_response['hits']['hits'].map { |result| result['_source'] }
+            current_results_length = current_results.length
+
+            # check if we reached the required offset
+            if results_offset < current_results_length
+              # check for parts
+              # (maybe a offset 6300 was provided but the batch size is 1000 - so we need to skip a part ...)
+              results_from = results_offset > 0 ? results_offset : 0
+              results_to   = (results_total + current_results_length - results_from) > results_limit ? results_limit - results_total + results_from - 1 : -1
+
+              ranged_results = current_results[results_from..results_to]
+
+              if block_given?
+                yield ranged_results
+              else
+                results |= ranged_results
+              end
+
+              # add to total
+              results_total += ranged_results.length
+            end
+
+            # -------- BREAK conditions --------
+
+            # we reached our maximum value
+            break if results_total >= results_limit
+
+            # we ran out of data
+            break if current_results_length < batch_size
+
+            # additional security - required?
+            # break if current_pit_hash[:search_after] == current_response['hits']['hits'][-1]['sort']
+
+            # -------- NEXT LOOP changes --------
+
+            # reduce the offset
+            results_offset -= current_results_length
+
+            # assign new pit
+            current_pit_hash = { search_after: current_response['hits']['hits'][-1]['sort'], pit: { id: current_response['pit_id'], keep_alive: keep_alive } }
+
+            # we need to justify the +batch_size+ if the query will reach over the limit
+            batch_size       = results_limit - results_total if results_offset < batch_size && (results_total + batch_size) > results_limit
+          end
+        end
+
+        # return results array
+        results
+      end
+
+      alias_method :total_results, :pit_results
+
+      # returns the RAW response for the current query
+      # @return [Array]
+      def response
+        spawn.hits_only!.resolve('Response').response
+      end
+
+      # returns the RAW aggregations for the current query
+      # @return [Hash]
+      def aggregations
+        spawn.aggs_only!.resolve('Aggregations').aggregations
+      end
+
+      # returns the RAW hits for the current query
+      # @return [Array]
+      def hits
+        spawn.hits_only!.resolve('Hits').hits
+      end
+
+      # returns the results for the current query
+      # @return [Array]
+      def results
+        spawn.hits_only!.resolve('Results').results
+      end
+
+      # returns the total value
+      def total
+        loaded? ? @total : spawn.total_only!.resolve('Total').total
+      end
+
+      # sets query as "hits"-only query (drops the aggs from the query)
+      def hits_only!
+        configure!({ aggs: nil })
+
+        self
+      end
+
+      # sets query as "aggs"-only query (drops the size & sort options - so no hits will return)
+      def aggs_only!
+        configure!({ size: 0, from: nil, sort: nil, _source: false })
+
+        self
+      end
+
+      # sets query as "total"-only query (drops the size, sort & aggs options - so no hits & aggs will be returned)
+      def total_only!
+        configure!({ size: 0, from: nil, aggs: nil, sort: nil, _source: false })
+
+        self
+      end
+    end
+  end
+end

data/lib/elasticsearch_record/relation/value_methods.rb

@@ -0,0 +1,54 @@
+module ElasticsearchRecord
+  module Relation
+    module ValueMethods
+      # holds the query kind
+      def kind_value
+        @values.fetch(:kind, nil)
+      end
+
+      def kind_value=(value)
+        # checks if records are already loaded - in this case we cannot mutate the query anymore
+        assert_mutability!
+
+        @values[:kind] = value.to_sym
+      end
+
+      def configure_value
+        @values.fetch(:configure, {})
+      end
+
+      def configure_value=(value)
+        assert_mutability!
+
+        @values[:configure] = value
+      end
+
+      def query_clause
+        @values.fetch(:query, ElasticsearchRecord::Relation::QueryClauseTree.empty)
+      end
+
+      def query_clause=(value)
+        assert_mutability!
+
+        @values[:query] = value
+      end
+
+      def aggs_clause
+        @values.fetch(:aggs, ElasticsearchRecord::Relation::QueryClauseTree.empty)
+      end
+
+      def aggs_clause=(value)
+        assert_mutability!
+
+        @values[:aggs] = value
+      end
+
+      private
+
+      # alternative method to avoid redefining the const +VALID_UNSCOPING_VALUES+
+      def _valid_unscoping_values
+        Set.new(ActiveRecord::Relation::VALID_UNSCOPING_VALUES.to_a + [:kind, :configure, :query, :aggs])
+      end
+    end
+  end
+end