2015 elasticsearch2 basic rails wrapper

searcheable.rb

module Searcheable
  extend ActiveSupport::Concern

  module ClassMethods

    def pluck(*fields)
      ElasticRelation.new(self).pluck(*fields)
    end

    def find_by(terms)
      where(terms).first
    end

    def first(n=1)
      ElasticRelation.new(self).first(n)
    end

    def last(n=1)
      ElasticRelation.new(self).last(n)
    end

    def limit(limit)
      ElasticRelation.new(self).limit(limit)
    end

    def sort(criteria)
      ElasticRelation.new(self).sort(criteria)
    end

    def count_all(terms)
      where(terms).count
    end

    def where_includes(terms, field='_all')
      # taks an array of terms to search for, and the field to search in
      ElasticRelation.new(self).where_includes(terms, field)
    end

    def where(terms)
      ElasticRelation.new(self).where(terms)
    end

    def where_not(terms)
      ElasticRelation.new(self).where_not(terms)
    end

    def where_not_exists(field)
      ElasticRelation.new(self).where_not_exists(field)
    end

    def where_exists(field)
      ElasticRelation.new(self).where_exists(field)
    end

    def custom(criteria)
      ElasticRelation.new(self).custom(criteria)
    end

    def cardinality(field)
      ElasticRelation.new(self).cardinality(field)
    end

    def agg(field, size=1000)
      ElasticRelation.new(self).agg(field, size)
    end

    def aggs(field1, field2, size=1000)
      ElasticRelation.new(self).aggs(field1, field2, size)
    end

    def destroy_all
      ElasticRelation.new(self).destroy_all
    end

  end

  class ElasticRelation
    def initialize(klass, query={'size'=> 1000})
      @klass = klass
      @query = query
    end

    ####################################################################################################################
    # Consolidating API
    ####################################################################################################################

    def all
      @klass.search(@query).results
    end

    def response
      @klass.search(@query).response
    end

    def hits
      response.hits
    end

    def first(n=1)
      n == 1 ? offset(0).limit(n).all[0] : offset(0).limit(n).all
    end

    def last(n=1)
      sort(created_at:'desc').first(n)
    end

    def count
      offset(0).limit(0).hits.total
    end

    def exists?
      count != 0
    end

    def cardinality(field)
      @query.merge!({'size'=>0, 'aggs' => { "distinct_field" => {'cardinality' => {'field' => field.to_s}}}})
      @klass.search(@query).response.aggregations.distinct_field.value
    end

    def agg(field, size=1000)
      @query.merge!({'size'=>0, 'aggs' => { "distinct_field" => {'terms' => {'field' => field.to_s, 'size' => size.to_s}}}})
      @klass.search(@query).response.aggregations.distinct_field.buckets
    end

    def aggs(field1, field2, size=1000)
      @query.merge!({'size'=>0, 'aggs' => { "distinct_field1" => {
          'terms' => {'field' => field1.to_s, 'size' => size.to_s},
          'aggs' => { "distinct_field2" => {
              'terms' => {'field' => field2.to_s, 'size' => size.to_s}
      }}}}})
      @klass.search(@query).response.aggregations.distinct_field1.buckets.inject({}){|s,z| s[z['key']]=z['distinct_field2']['buckets'].inject({}){|q,w| q[w['key']]=w['doc_count'];q}; s}
    end

    def pluck(*fields)
      fields.map!(&:to_s)
      @query.merge!({'fields': fields})
      response = @klass.search(@query).response
      hits = response.hits.hits.map{|h| (h.fields || {}).merge({"id"=>h._id})}
      if fields.size == 1
        hits.flat_map{|hit| fields.flat_map{|field| hit[field] unless hit.nil?}}
      else
        hits.map{|hit| fields.flat_map{|field| hit[field] unless hit.nil?}}
      end
    end

    def destroy_all
      @klass.find_in_batches(@query) do |batch|
        commands = batch.map do |register|
          {"delete":{"_index":"#{@klass.index_name}","_type":"#{@klass.document_type}", "_id":"#{register.id}"}}
        end
        Elasticsearch::Persistence.client.bulk({body:commands})
      end
    end

    def to_sql
      @query
    end

    ####################################################################################################################
    # transforming API
    ####################################################################################################################

    def validate_terms(terms)
      if !terms || terms == {}
        return false
      end
      true
    end

    def where_includes(terms, field='_all')
      return self unless validate_terms(terms)
      initialize_query_string(field)
      @query['query']['filtered']['query']['query_string'] = mount_where_includes_criteria(terms, field)
      self
    end

    def where(terms)
      return self unless validate_terms(terms)
      terms.stringify_keys!
      initialize_filter('must')
      @query['query']['filtered']['filter']['bool']['must'] << mount_where_criteria(terms)
      self
    end

    def where_not(terms)
      return self unless validate_terms(terms)
      terms.stringify_keys!
      initialize_filter('must_not')
      @query['query']['filtered']['filter']['bool']['must_not'] << mount_where_criteria(terms)
      self
    end

    def where_exists(field)
      return self unless validate_terms(field)
      initialize_filter('must')
      @query['query']['filtered']['filter']['bool']['must'] << {'exists'=>{'field'=>field}}
      self
    end

    def where_not_exists(field)
      initialize_filter('must_not')
      @query['query']['filtered']['filter']['bool']['must_not'] << {'exists'=>{'field'=>field.to_s}}
      self
    end

    def offset(size)
      @query.deep_merge!({'from' => size})
      self
    end

    def limit(size)
      @query.deep_merge!({'size' => size})
      self
    end

    def sort(criteria)
      @query.deep_merge!({'sort'=> mount_sort_criteria(criteria)})
      self
    end

    def custom(criteria)
      @query.deep_merge!(criteria)
      self
    end

    private

    def mount_sort_criteria(criteria)
      return mount_sort_criteria(criteria => 'asc') if criteria.is_a?(String) || criteria.is_a?(Symbol)
      criteria.stringify_keys!
      criteria.map{|k,v| {k=>{'order' => v}}}
    end

    def mount_where_includes_criteria(values, field)
      values = [values] if values.is_a?(String)
      values = values.map do |f| "*#{f}*" end
      values = values.join(' OR ')
      { "default_field" => field, "query" => values }
    end

    def mount_where_criteria(criteria)
      criteria.map do |k,v|
        case v
          when Array then {'terms'=>{k=>v}}
          when Range then {'range' => {k => {'from'=> v.min, 'to' => v.max}}}
          else {'term'=>{k=>v}}
        end
      end
    end

    def initialize_filter(filter)
      @query = {'query'=>{'filtered'=>{'filter'=>{'bool'=>{filter=>[]}}}}}.deep_merge(@query)
    end

    def initialize_query_string(filter)
      @query = {'query'=>{"filtered"=>{"query"=>{'query_string'=>{}}}}}.deep_merge(@query)
    end
  end
end