quandl-elasticsearch 2.1.0.rc5

data/lib/quandl/elasticsearch/query.rb

@@ -0,0 +1,282 @@
+module Quandl
+  module Elasticsearch
+    class Query # rubocop:disable Metrics/ClassLength
+      class << self
+        SAMPLE_FILTER = 'sample'.freeze
+
+        def add_timeout!(body)
+          # see: https://www.elastic.co/guide/en/elasticsearch/reference/current/search-request-body.html
+          api_timeout = Quandl::Elasticsearch.configuration.api_timeout
+          return unless api_timeout
+
+          timeout_val = "#{api_timeout}s"
+          if body[:search]
+            body[:search][:timeout] = timeout_val
+          else
+            body[:timeout] = timeout_val
+          end
+        end
+
+        def query_body(query, page, per_page)
+          default_should = [multi_match(query), child?(query)]
+          should = single_term?(query) ? default_should << prefix(query) : default_should
+          default_query = {
+            bool: {
+              must: {
+                bool: {
+                  should: should
+                }
+              },
+              filter: {
+                bool: {
+                  must: []
+                }
+              }
+            }
+          }
+
+          body = {
+            explain: false,
+            fields: %w(_parent code),
+            query: try_match_all_query(query, default_query),
+            from: page, size: per_page
+          }
+
+          body[:rescore] = rescore(query) if page == 0
+          body[:sort] = [{ rating: { order: 'desc' } }] if query.empty?
+          body
+        end
+
+        def try_match_all_query(query, default_query)
+          match_all_query = {
+            bool: {
+              must: { match_all: {} },
+              filter: {
+                bool: {
+                  must: []
+                }
+              }
+            }
+          }
+          query.to_s.empty? ? match_all_query : default_query
+        end
+
+        def multi_dataset_query_body_should(query)
+          should = [
+            {
+              multi_match: {
+                fields: ['code', 'name^1.1'],
+                query: query.to_s,
+                type: 'best_fields'
+              }
+            }
+          ]
+          should << { prefix: { code: query.to_s } } if single_term?(query)
+          should
+        end
+
+        def multi_dataset_query_body(id, query, max_datasets)
+          should = multi_dataset_query_body_should(query)
+
+          body = {
+            index: 'quandl_index', type: 'dataset',
+            search: {
+              fields: %w(_parent _id code),
+              query:         {
+                bool: {
+                  must: {
+                    bool: {
+                      should: should
+                    }
+                  },
+                  filter: {
+                    bool: {
+                      must: [
+                        { term: { _parent: id } },
+                        { term: { is_private: false } }
+                      ]
+                    }
+                  }
+                }
+              },
+              sort: ['_score', { to_date: { order: 'desc' } }],
+              highlight: { fields: { name: {} } },
+              size: max_datasets
+            }
+          }
+
+          add_timeout!(body)
+          body
+        end
+
+        def multi_dataset_query_body_match_all(id, max_datasets)
+          body = {
+            index: 'quandl_index', type: 'dataset',
+            search: {
+              fields: %w(_parent _id code),
+              query:             {
+                bool: {
+                  must: {
+                    match_all: {}
+                  },
+                  filter: {
+                    bool: {
+                      must: [
+                        { term: { _parent: id } },
+                        { term: { is_private: false } }
+                      ]
+                    }
+                  }
+                }
+              },
+              sort: [{ _id: { order: 'desc' } }],
+              highlight: { fields: { name: {} } },
+              size: max_datasets
+            }
+          }
+
+          add_timeout!(body)
+          body
+        end
+
+        def dataset_query_body(from, size, options = {})
+          query_body = {
+            fields: %w(_parent code),
+            query: {
+              bool: {
+                filter: {
+                  bool: {
+                    # Filter out any datasets that are private
+                    must: [{ term: { is_private: false } }]
+                  }
+                }
+              }
+            },
+            from: from,
+            size: size
+          }
+
+          # Search is filtering by database_id
+          if options[:database_id]
+            query_body[:query][:bool][:filter][:bool][:must] << { term: { _parent: options[:database_id] } }
+          else
+            # if we are not filtering by database_id, ensure that the datasets we return do not have a hidden database
+            # only return datasets that are NOT exclusive
+            query_body[:query][:bool][:filter][:bool][:must] += [{ term: { db_hidden: false } }, { term: { db_exclusive: false } }]
+          end
+
+          # add filters specified by request
+          if options[:filter] == SAMPLE_FILTER
+            query_body[:query][:bool][:filter][:bool][:must] << { term: { sample: true } }
+          end
+
+          # Searching with a query string or without
+          if options[:query] && !options[:query].empty?
+            query_body[:query][:bool][:must] = { bool: { should: should(options[:query]) } }
+            query_body[:highlight] = { fields: { name: {} } }
+          else
+            query_body[:sort] = [{ to_date: { order: 'desc' } }]
+            query_body[:query][:bool][:must] = { match_all: {} }
+          end
+
+          # Sort out the search results for consistency
+          case options[:sort_by]
+          when 'id'
+            query_body[:sort] = [{ _uid: { order: 'asc' } }]
+          end
+
+          query_body
+        end
+
+        def term_tags(tags)
+          tags.map do |tag|
+            { term: { tags: tag } }
+          end
+        end
+
+        private
+
+        def single_term?(query)
+          query && query.split(' ').length == 1
+        end
+
+        def multi_match(query)
+          default_fields = ['name^1.1']
+          fields = single_term?(query) ? ['code'] + default_fields : default_fields
+          {
+            multi_match: {
+              fields: fields,
+              query: query.to_s,
+              type: 'phrase',
+              operator: 'and',
+              zero_terms_query: 'all',
+              slop: 10
+            }
+          }
+        end
+
+        def child?(query)
+          should = should(query)
+          {
+            has_child: {
+              type: 'dataset',
+              score_mode: 'max',
+              query: {
+                bool: {
+                  should: should
+                }
+              }
+            }
+          }
+        end
+
+        # https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-multi-match-query.html#type-cross-fields
+        def should(query)
+          default_fields = ['name^1.1']
+          fields = single_term?(query) ? ['code'] + default_fields : default_fields
+          default_should = [
+            {
+              multi_match: {
+                fields: default_fields,
+                query: query.to_s,
+                type: 'phrase',
+                slop: 10
+              }
+            },
+            {
+              multi_match: {
+                fields: fields,
+                query: query.to_s,
+                type: 'best_fields'
+              }
+            }
+          ]
+          single_term?(query) ? default_should << { prefix: { code: query.to_s } } : default_should
+        end
+
+        def rescore(query)
+          {
+            window_size: 10,
+            query: {
+              rescore_query: {
+                match: {
+                  name: {
+                    query: query.to_s,
+                    type: 'phrase',
+                    slop: 10
+                  }
+                }
+              },
+              query_weight: 0.7,
+              rescore_query_weight: 1.2
+            }
+          }
+        end
+
+        def prefix(query)
+          { prefix: { code: query.to_s } }
+        end
+      end
+    end
+  end
+end

data/lib/quandl/elasticsearch/search.rb

@@ -0,0 +1,150 @@
+require 'quandl/elasticsearch/tag'
+require 'quandl/elasticsearch/query'
+
+module Quandl
+  module Elasticsearch
+    class Search
+      class TimeoutError < StandardError; end
+
+      attr_accessor :page_size
+      PER_PAGE = 10
+      MAX_DATASETS = 3
+      DEFAULT_SCOPE = 'all'
+
+      def initialize(page_size = PER_PAGE)
+        @client = Quandl::Elasticsearch.client
+        @page_size = [1000, page_size].min
+      end
+
+      def page_size
+        @page_size ||= PER_PAGE
+      end
+
+      def max_datasets
+        @max_datasets ||= MAX_DATASETS
+      end
+
+      def database(query, tags = '', scope = DEFAULT_SCOPE, page = 0, with_dataset = true)
+        query = query.to_s unless query.is_a? String
+        query.downcase!
+        return_premium, return_all = interpret_scope(scope)
+        tags = tags.split(',')
+        query = normalize_query(query)
+        from = page * @page_size
+        body = Quandl::Elasticsearch::Query.query_body(query, from, page_size)
+
+        body[:query][:bool][:filter][:bool][:must] << { term: { hidden: false } }
+        # only include databases that are not exclusive
+        body[:query][:bool][:filter][:bool][:must] << { term: { exclusive: false } }
+        body[:query][:bool][:filter][:bool][:must] += Quandl::Elasticsearch::Query.term_tags(tags)
+        body[:query][:bool][:filter][:bool][:must] << { term: { premium: return_premium } } unless return_all
+
+        result = es_query('quandl_index', 'database', body)
+        result_set = result['hits']['hits'].map { |r| { '_id' => r['_id'], 'code' => r['fields']['code'].first } }
+
+        if with_dataset
+          database_ids = result_set.map { |d| d['_id'] }
+          datasets = query_dataset(database_ids, query)
+          result_set.each { |d| d['datasets'] = datasets["#{d['_id']}"] ? datasets["#{d['_id']}"] : [] }
+        end
+
+        if datasets.nil? || datasets && datasets.empty?
+          took = result['took']
+        else
+          took = result['took'] + datasets['took']
+        end
+        result_set << { 'total' => result['hits']['total'], 'took' =>  took }
+        result_set
+      end
+
+      def dataset(query, frequency, page = 0, size = page_size, options = {})
+        from = page * size
+        query = normalize_query(query)
+
+        options[:query] = query unless query =~ /^\s+$/ || query == ''
+
+        body = Quandl::Elasticsearch::Query.dataset_query_body(from, size, options)
+
+        unless frequency.nil?
+          body[:query][:bool][:filter][:bool][:should] = []
+          frequency.split(',').each do |f|
+            body[:query][:bool][:filter][:bool][:should] << { term: { frequency: "#{f}" } }
+          end
+        end
+        result = es_query('quandl_index', ['dataset'], body)
+        convert_dataset_result(result)
+      end
+
+      private
+
+      def es_query(index, type, body)
+        Quandl::Elasticsearch::Query.add_timeout!(body)
+        result = @client.search index: index, type: type, body: body
+        handle_timeout!(result)
+        result
+      end
+
+      def handle_timeout!(result)
+        raise Quandl::Elasticsearch::Search::TimeoutError if result['timed_out'] # rubocop:disable Style/SignalException
+      end
+
+      def query_dataset(database_ids, query)
+        return [] if database_ids.empty?
+        body = []
+        database_ids.each { |id| body << Quandl::Elasticsearch::Query.multi_dataset_query_body(id, query, max_datasets) } unless query.to_s.empty?
+        database_ids.each { |id| body << Quandl::Elasticsearch::Query.multi_dataset_query_body_match_all(id, max_datasets) } if query.to_s.empty?
+
+        r = @client.msearch body: body
+        result = {}
+        r['responses'].each do |ds|
+          handle_timeout!(ds)
+          next if ds['hits']['hits'].empty?
+          database_id = ds['hits']['hits'].first['_parent']
+          result[database_id] = ds['hits']['hits'].map do |d|
+            { 'id' => d['_id'],
+              'code' => d['fields']['code'].first,
+              'highlight' => d['highlight'] }
+          end
+        end
+        result['took'] = r['responses'].map { |record| record['took'] }.max
+        result
+      end
+
+      private
+
+      # Ensure query is a string and downcased and not a number
+      def normalize_query(query)
+        query.to_s.downcase
+      end
+
+      def convert_dataset_result(result)
+        result_set = result['hits']['hits'].map do |r|
+          { '_id' => r['_id'],
+            'highlight' => r['highlight']
+            }
+        end
+        result_set << { 'total' => result['hits']['total'] }
+        result_set
+      end
+
+      def interpret_scope(scope = DEFAULT_SCOPE)
+        return true, false if scope.class == TrueClass # support command line tool for now
+        return false, false if scope.class == FalseClass # support command line tool for now
+        scope = DEFAULT_SCOPE if scope.nil?
+        scope.upcase!
+        premium = scope == 'PREMIUM' ? true : false
+        all = scope == 'ALL' ? true : false
+        # rubocop:disable RedundantReturn
+        return premium, all
+        # rubocop:enable RedundantReturn
+      end
+
+      def tag_include?(a1, a2)
+        a2.each do |a|
+          return true if a1.include? a
+        end
+        false
+      end
+    end
+  end
+end

data/lib/quandl/elasticsearch/tag.rb

@@ -0,0 +1,21 @@
+module Quandl
+  module Elasticsearch
+    class Tag
+      def initialize(tags)
+        @tags = tags.to_s
+      end
+
+      def to_s
+        @tags
+      end
+
+      def to_a
+        @tags.split(',')
+      end
+
+      def blank?
+        @tags.empty?
+      end
+    end
+  end
+end