embulk-input-elasticsearch 0.2.1 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 24e8943311e14d4c3ba8fbbdb490bed67a2974ba
-  data.tar.gz: 4f795401dfb26d3e5b58e39d4921a99eb1530445
+  metadata.gz: 0991bd544cf235c70290b6d307f164112280d1ea
+  data.tar.gz: 4e578ad476510b38eaa09ef9a20437f1548cc70c
 SHA512:
-  metadata.gz: 0b48169a99f7d055271aeecf213704bd1798d9b4f0bac14b28a56933c5ca3ce0bce48b849690abe8a133b7f2e06d4bd869c2cdb01351904fa9992bc7f429efa8
-  data.tar.gz: 1d2dafc4528c3d306f7bbbd08de973b3bdc2ec0ab4dc6ae0997aeb1adbe83e940d71e9becf0c9dcafeebf8169a76275fb9ce3bb4c0e3a6d9dd1a9d3ad4612a1d
+  metadata.gz: b47b9034846fc515fda638e231b5c62f13b0c8fecf1fba8ef16ad8f084f967518b2089df5a95e7681bb6e9af2dd918da452b843d1072dd04e22f01f504af674e
+  data.tar.gz: ddc4879a8a6c0524499f1b10f3b1e871c1920024aefa81e1c4bf562c44533cc161ea6d8d202bf1dd89af30688f35790c3f2f9b0a9a637a9af9ebdd86e32b6e09

data/embulk-input-elasticsearch.gemspec

@@ -1,7 +1,7 @@
 
 Gem::Specification.new do |spec|
   spec.name          = "embulk-input-elasticsearch"
-  spec.version       = "0.2.1"
+  spec.version       = "0.3.0"
   spec.authors       = ["toyama0919"]
   spec.summary       = "Elasticsearch input plugin for Embulk"
   spec.description   = "Loads records from Elasticsearch. parallel query support."

data/lib/embulk/input/elasticsearch.rb

@@ -1,11 +1,14 @@
 require 'excon'
 require 'elasticsearch'
+require_relative 'elasticsearch/connection'
+require_relative 'elasticsearch/input_thread'
 
 module Embulk
   module Input
 
     class Elasticsearch < InputPlugin
       Plugin.register_input("elasticsearch", self)
+      ADD_QUERY_TO_RECORD_KEY = 'query'
 
       def self.transaction(config, &control)
         task = {
@@ -25,31 +28,19 @@ module Embulk
         }
         # TODO: want max_threads
         define_num_threads = config.param("num_threads", :integer, default: 1)
-        task['slice_queries'] = get_slice_from_num_threads(task['queries'], define_num_threads)
+        task['slice_queries'] = InputThread.get_slice_from_num_threads(task['queries'], define_num_threads)
 
         columns = []
         task['fields'].each_with_index{ |field, i|
           columns << Column.new(i, field['name'], field['type'].to_sym)
         }
         if task['add_query_to_record']
-          columns << Column.new(task['fields'].size, "query", :string)
+          columns << Column.new(task['fields'].size, ADD_QUERY_TO_RECORD_KEY, :string)
         end
 
         resume(task, columns, task['slice_queries'].size, &control)
       end
 
-      def self.get_slice_from_num_threads(array, define_num_threads)
-        num_threads = array.size < define_num_threads ? array.size : define_num_threads
-        per_queries = if (array.size % num_threads) == 0
-          (array.size / num_threads)
-        else
-          (array.size / num_threads) + 1
-        end
-        sliced = array.each_slice(per_queries).to_a
-        Embulk.logger.info("calculate num threads => #{sliced.size}")
-        return sliced
-      end
-
       def self.resume(task, columns, count, &control)
         task_reports = yield(task, columns, count)
 
@@ -57,28 +48,10 @@ module Embulk
         return next_config_diff
       end
 
-      def self.create_client(task)
-        transport = ::Elasticsearch::Transport::Transport::HTTP::Faraday.new(
-          {
-            hosts: task['nodes'].map{ |node| Hash[node.map{ |k, v| [k.to_sym, v] }] },
-            options: {
-              reload_connections: task['reload_connections'],
-              reload_on_failure: task['reload_on_failure'],
-              retry_on_failure: task['retry_on_failure'],
-              transport_options: {
-                request: { timeout: task['request_timeout'] }
-              }
-            }
-          }
-        )
-
-        ::Elasticsearch::Client.new transport: transport
-      end
-
       def init
         @queries = task['slice_queries'][@index]
         Embulk.logger.info("this thread queries => #{@queries}")
-        @client = self.class.create_client(task)
+        @client = Connection.create_client(task)
         @index_name = task['index']
         @index_type = task['index_type']
         @per_size = task['per_size']
@@ -89,27 +62,7 @@ module Embulk
       end
 
       def run
-        @queries.each do |query|
-          query_count = 0
-          no_source_results = search(@index_type, query, 0, 0, @routing, @fields, @sort)
-          total_count = [no_source_results['hits']['total'], @limit_size].compact.min
-          while true
-            now_results_size = query_count * @per_size
-            next_results_size = (query_count + 1) * @per_size
-            size = get_size(next_results_size, now_results_size ,total_count)
-            break if size == 0
-
-            results = get_sources(search(@index_type, query, size, now_results_size, @routing, @fields, @sort), @fields)
-            results.each do |record|
-              if @add_query_to_record
-                record << query
-              end
-              page_builder.add(record)
-            end
-            break if last_query?(next_results_size ,total_count)
-            query_count += 1
-          end
-        end
+        search(@index_type, @per_size, @routing, @fields, @sort)
         page_builder.finish
 
         task_report = {}
@@ -118,65 +71,52 @@ module Embulk
 
       private
 
-      def convert_value(value, field)
-        return nil if value.nil?
-        case field["type"]
-        when "string"
-          value
-        when "long"
-          value.to_i
-        when "double"
-          value.to_f
-        when "boolean"
-          if value.is_a?(TrueClass) || value.is_a?(FalseClass)
-            value
-          else
-            downcased_val = value.downcase
-            case downcased_val
-            when 'true' then true
-            when 'false' then false
-            when '1' then true
-            when '0' then false
-            else nil
-            end
-          end
-        when "timestamp"
-          Time.parse(value)
-        when "json"
-          value
-        else
-          raise "Unsupported type #{field['type']}"
+      def search(type, size, routing, fields, sort)
+        @queries.each do |query|
+          search_with_query(query, type, size, routing, fields, sort)
         end
       end
 
-      def get_size(next_results_size, now_results_size ,total_count)
-        if last_query?(next_results_size ,total_count)
-          (total_count - now_results_size)
-        else
-          @per_size
+      def search_with_query(query, type, size, routing, fields, sort)
+        search_option = get_search_option(type, query, size, fields, sort)
+        Embulk.logger.info("#{search_option}")
+        r = @client.search(search_option)
+        i = 0
+        get_sources(r, fields).each do |result|
+          result_proc(result, query)
+          return if @limit_size == (i += 1)
+        end
+
+        while r = @client.scroll(scroll_id: r['_scroll_id'], scroll: '1m') and (not r['hits']['hits'].empty?) do
+          get_sources(r, fields).each do |result|
+            result_proc(result, query)
+            return if @limit_size == (i += 1)
+          end
         end
       end
 
-      def last_query?(next_results_size ,total_count)
-        next_results_size > total_count
+      def result_proc(result, query)
+        if @add_query_to_record
+          result << query
+        end
+        page_builder.add(result)
       end
 
-      def search(type, query, size, from, routing, fields, sort)
-        body = { from: from }
-        body[:size] = size unless size.nil?
+      def get_search_option(type, query, size, fields, sort)
+        body = { }
+        body[:query] = { query_string: { query: query } } unless query.nil?
         if sort
           sorts = []
           sort.each do |k, v|
             sorts << { k => v }
           end
           body[:sort] = sorts
+        else
+          body[:sort] = ["_doc"]
         end
-        body[:query] = { query_string: { query: query } } unless query.nil?
-        search_option = { index: @index_name, type: type, body: body }
-        search_option[:routing] = routing unless routing.nil?
+        search_option = { index: @index_name, type: type, scroll: '1m', body: body, size: size }
         search_option[:_source] = fields.select{ |field| !field['metadata'] }.map { |field| field['name'] }.join(',')
-        Embulk.logger.info(%Q{search_option => #{search_option}})
-        @client.search(search_option)
+        search_option
       end
 
       def get_sources(results, fields)
@@ -191,6 +131,37 @@ module Embulk
           }
         }
       end
+
+      def convert_value(value, field)
+        return nil if value.nil?
+        case field["type"]
+        when "string"
+          value
+        when "long"
+          value.to_i
+        when "double"
+          value.to_f
+        when "boolean"
+          if value.is_a?(TrueClass) || value.is_a?(FalseClass)
+            value
+          else
+            downcased_val = value.downcase
+            case downcased_val
+            when 'true' then true
+            when 'false' then false
+            when '1' then true
+            when '0' then false
+            else nil
+            end
+          end
+        when "timestamp"
+          Time.parse(value)
+        when "json"
+          value
+        else
+          raise "Unsupported type #{field['type']}"
+        end
+      end
     end
   end
 end

data/lib/embulk/input/elasticsearch/connection.rb

@@ -0,0 +1,25 @@
+module Embulk
+  module Input
+    class Elasticsearch < InputPlugin
+      class Connection
+        def self.create_client(task)
+          transport = ::Elasticsearch::Transport::Transport::HTTP::Faraday.new(
+            {
+              hosts: task['nodes'].map{ |node| Hash[node.map{ |k, v| [k.to_sym, v] }] },
+              options: {
+                reload_connections: task['reload_connections'],
+                reload_on_failure: task['reload_on_failure'],
+                retry_on_failure: task['retry_on_failure'],
+                transport_options: {
+                  request: { timeout: task['request_timeout'] }
+                }
+              }
+            }
+          )
+
+          ::Elasticsearch::Client.new transport: transport
+        end
+      end
+    end
+  end
+end

data/lib/embulk/input/elasticsearch/input_thread.rb

@@ -0,0 +1,19 @@
+module Embulk
+  module Input
+    class Elasticsearch < InputPlugin
+      class InputThread
+        def self.get_slice_from_num_threads(array, define_num_threads)
+          num_threads = array.size < define_num_threads ? array.size : define_num_threads
+          per_queries = if (array.size % num_threads) == 0
+            (array.size / num_threads)
+          else
+            (array.size / num_threads) + 1
+          end
+          sliced = array.each_slice(per_queries).to_a
+          Embulk.logger.info("calculate num threads => #{sliced.size}")
+          return sliced
+        end
+      end
+    end
+  end
+end

data/test/test_transaction.rb

@@ -13,25 +13,25 @@ module Embulk
 
       sub_test_case "get_slice_from_num_threads" do
         def test_normal
-          slice = Elasticsearch.get_slice_from_num_threads((1..10).to_a, 5)
+          slice = InputThread.get_slice_from_num_threads((1..10).to_a, 5)
           assert_equal slice.size, 5
           assert_equal slice.first.size, 2
         end
 
         def test_normal_same
-          slice = Elasticsearch.get_slice_from_num_threads((1..3).to_a, 3)
+          slice = InputThread.get_slice_from_num_threads((1..3).to_a, 3)
           assert_equal slice.size, 3
           assert_equal slice.first.size, 1
         end
 
         def test_num_threads_over_array_size
-          slice = Elasticsearch.get_slice_from_num_threads((1..3).to_a, 10)
+          slice = InputThread.get_slice_from_num_threads((1..3).to_a, 10)
           assert_equal slice.size, 3
           assert_equal slice.first.size, 1
         end
 
         def test_rest
-          slice = Elasticsearch.get_slice_from_num_threads((1..20).to_a, 8)
+          slice = InputThread.get_slice_from_num_threads((1..20).to_a, 8)
           assert_equal slice.size, 7
           assert_equal slice.first.size, 3
         end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: embulk-input-elasticsearch
 version: !ruby/object:Gem::Version
-  version: 0.2.1
+  version: 0.3.0
 platform: ruby
 authors:
 - toyama0919
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2016-06-21 00:00:00.000000000 Z
+date: 2016-06-24 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   requirement: !ruby/object:Gem::Requirement
@@ -110,6 +110,8 @@ files:
 - Rakefile
 - embulk-input-elasticsearch.gemspec
 - lib/embulk/input/elasticsearch.rb
+- lib/embulk/input/elasticsearch/connection.rb
+- lib/embulk/input/elasticsearch/input_thread.rb
 - test/helper.rb
 - test/test_transaction.rb
 homepage: https://github.com/toyama0919/embulk-input-elasticsearch