chewy 0.9.0 → 0.10.0

data/lib/chewy/type/adapter/orm.rb

@@ -6,16 +6,15 @@ module Chewy
       class Orm < Base
         attr_reader :default_scope
 
-        def initialize(*args)
-          @options = args.extract_options!
-          class_or_relation = args.first
-          if class_or_relation.is_a?(relation_class)
-            @target = model_of_relation(class_or_relation)
-            @default_scope = class_or_relation
+        def initialize(target, **options)
+          if target.is_a?(relation_class)
+            @target = model_of_relation(target)
+            @default_scope = target
           else
-            @target = class_or_relation
+            @target = target
             @default_scope = all_scope
           end
+          @options = options
           cleanup_default_scope!
         end
 
@@ -25,10 +24,10 @@ module Chewy
 
         def identify(collection)
           if collection.is_a?(relation_class)
-            pluck_ids(collection)
+            pluck(collection)
           else
             Array.wrap(collection).map do |entity|
-              entity.is_a?(object_class) ? entity.public_send(primary_key) : entity
+              entity.respond_to?(primary_key) ? entity.public_send(primary_key) : entity
             end
           end
         end
@@ -49,7 +48,7 @@ module Chewy
         #
         # Method handles destroyed objects as well. In case of objects ORM scope
         # or array passed, objects, responding with true to `destroyed?` method will be deleted
-        # from index. In case of ids array passed - documents with missing records ids will be
+        # from index. In case of ids array passed - documents with missing source object ids will be
         # deleted from index:
         #
         #   users = User.all
@@ -74,16 +73,7 @@ module Chewy
         #   UsersIndex::User.import users.map(&:id) # user ids will be deleted from index
         #
         def import(*args, &block)
-          options = args.extract_options!
-          options[:batch_size] ||= BATCH_SIZE
-
-          collection = if args.empty?
-            default_scope
-          elsif args.one? && args.first.is_a?(relation_class)
-            args.first
-          else
-            args.flatten.compact
-          end
+          collection, options = import_args(*args)
 
           if collection.is_a?(relation_class)
             import_scope(collection, options, &block)
@@ -92,29 +82,51 @@ module Chewy
           end
         end
 
-        def load(*args)
-          load_options = args.extract_options!
-          objects = args.flatten
+        def import_fields(*args, &block)
+          return enum_for(:import_fields, *args) unless block_given?
+
+          collection, options = import_args(*args)
+
+          if options[:fields].present? || collection.is_a?(relation_class)
+            collection = all_scope_where_ids_in(identify(collection)) unless collection.is_a?(relation_class)
+            pluck_in_batches(collection, options.slice(:fields, :batch_size, :typecast), &block)
+          else
+            identify(collection).each_slice(options[:batch_size]) do |batch|
+              yield batch
+            end
+          end
+        end
+        alias_method :import_references, :import_fields
 
-          additional_scope = load_options[load_options[:_type].type_name.to_sym].try(:[], :scope) || load_options[:scope]
+        def load(ids, **options)
+          scope = all_scope_where_ids_in(ids)
+          additional_scope = options[options[:_type].type_name.to_sym].try(:[], :scope) || options[:scope]
 
-          scope = all_scope_where_ids_in(objects.map(&primary_key))
-          loaded_objects = load_scope_objects(scope, additional_scope).index_by { |object| object.public_send(primary_key).to_s }
+          loaded_objects = load_scope_objects(scope, additional_scope)
+            .index_by do |object|
+              object.public_send(primary_key).to_s
+            end
 
-          objects.map { |object| loaded_objects[object.public_send(primary_key).to_s] }
+          ids.map { |id| loaded_objects[id.to_s] }
         end
 
       private
 
         def import_objects(collection, options)
-          hash = Hash[identify(collection).zip(collection)]
+          collection_ids = identify(collection)
+          hash = Hash[collection_ids.map(&:to_s).zip(collection)]
+
+          indexed = collection_ids.each_slice(options[:batch_size]).map do |ids|
+            batch = if options[:raw_import]
+              raw_default_scope_where_ids_in(ids, options[:raw_import])
+            else
+              default_scope_where_ids_in(ids)
+            end
 
-          indexed = hash.keys.each_slice(options[:batch_size]).map do |ids|
-            batch = default_scope_where_ids_in(ids)
             if batch.empty?
               true
             else
-              identify(batch).each { |id| hash.delete(id) }
+              batch.each { |object| hash.delete(object.send(primary_key).to_s) }
               yield grouped_objects(batch)
             end
           end.all?
@@ -157,7 +169,22 @@ module Chewy
         end
 
         def grouped_objects(objects)
-          options[:delete_if] ? super : { index: objects.to_a }
+          options[:delete_if] ? super : {index: objects.to_a}
+        end
+
+        def import_args(*args)
+          options = args.extract_options!
+          options[:batch_size] ||= BATCH_SIZE
+
+          collection = if args.empty?
+            default_scope
+          elsif args.one? && args.first.is_a?(relation_class)
+            args.first
+          else
+            args.flatten.compact
+          end
+
+          [collection, options]
         end
       end
     end

data/lib/chewy/type/adapter/sequel.rb

@@ -24,38 +24,54 @@ module Chewy
         end
 
         def import_scope(scope, options)
-          scope = scope.unordered.order(::Sequel.asc(primary_key_with_table_name)).limit(options[:batch_size])
-
-          ids = pluck_ids(scope)
-          result = true
-
-          while ids.present?
-            result &= yield grouped_objects(default_scope_where_ids_in(ids).all)
-            break if ids.size < options[:batch_size]
-            ids = pluck_ids(scope.where { |o| o.__send__(primary_key_with_table_name) > ids.last })
+          pluck_in_batches(scope, options.slice(:batch_size)).inject(true) do |result, ids|
+            result & yield(grouped_objects(default_scope_where_ids_in(ids).all))
           end
-
-          result
         end
 
         def primary_key
           target.primary_key
         end
 
-        def primary_key_with_table_name
-          "#{target.table_name}__#{primary_key}".to_sym
+        def full_column_name(column)
+          "#{target.table_name}__#{column}".to_sym
         end
 
         def all_scope
           target.dataset
         end
 
-        def pluck_ids(scope)
-          scope.distinct.select_map(primary_key_with_table_name)
+        def target_columns
+          @target_columns ||= target.columns.to_set
+        end
+
+        def pluck(scope, fields: [])
+          fields = fields.map(&:to_sym).unshift(primary_key).map do |column|
+            target_columns.include?(column) ? full_column_name(column) : column
+          end
+          scope.distinct.select_map(fields.one? ? fields.first : fields)
+        end
+
+        def pluck_in_batches(scope, fields: [], batch_size: nil, **options)
+          return enum_for(:pluck_in_batches, scope, fields: fields, batch_size: batch_size, **options) unless block_given?
+
+          scope = scope.unordered.order(::Sequel.asc(full_column_name(primary_key))).limit(batch_size)
+
+          ids = pluck(scope, fields: fields)
+          count = 0
+
+          while ids.present?
+            yield ids
+            break if ids.size < batch_size
+            last_id = ids.last.is_a?(Array) ? ids.last.first : ids.last
+            ids = pluck(scope.where { |o| o.__send__(full_column_name(primary_key)) > last_id }, fields: fields)
+          end
+
+          count
         end
 
         def scope_where_ids_in(scope, ids)
-          scope.where(primary_key_with_table_name => Array.wrap(ids))
+          scope.where(full_column_name(primary_key) => Array.wrap(ids))
         end
 
         def model_of_relation(relation)

data/lib/chewy/type/import.rb

@@ -1,241 +1,195 @@
+require 'chewy/type/import/journal_builder'
+require 'chewy/type/import/bulk_builder'
+require 'chewy/type/import/bulk_request'
+require 'chewy/type/import/routine'
+
 module Chewy
   class Type
     module Import
       extend ActiveSupport::Concern
 
-      BULK_OPTIONS = [:suffix, :bulk_size, :refresh, :consistency, :replication].freeze
+      IMPORT_WORKER = lambda do |type, options, ids|
+        ::Process.setproctitle("chewy import #{type}[#{::Parallel.worker_number}]")
+        routine = Routine.new(type, options)
+        type.adapter.import(*ids, routine.options) do |action_objects|
+          routine.process(**action_objects)
+        end
+        {errors: routine.errors, import: routine.stats, leftovers: routine.leftovers}
+      end
+
+      LEFTOVERS_WORKER = lambda do |type, options, body|
+        ::Process.setproctitle("chewy import #{type}[#{::Parallel.worker_number}]")
+        routine = Routine.new(type, options)
+        routine.perform_bulk(body)
+        routine.errors
+      end
 
       module ClassMethods
-        # Perform import operation for specified documents.
-        # Returns true or false depending on success.
+        # @!method import(*collection, **options)
+        # Basically, one of the main methods for type. Performs any objects import
+        # to the index for a specified type. Does all the objects handling routines.
+        # Performs document import by utilizing bulk API. Bulk size and objects batch
+        # size are controlled by the corresponding options.
+        #
+        # It accepts ORM/ODM objects, PORO, hashes, ids which are used by adapter to
+        # fetch objects from the source depenting on the used adapter. It destroys
+        # passed objects from the index if they are not in the default type scope
+        # or marked for destruction.
+        #
+        # It handles parent-child relationships: if the object parent_id has been
+        # changed it destroys the object and recreates it from scratch.
+        #
+        # Performs journaling if enabled: it stores all the ids of the imported
+        # objects to a specialized index. It is possible to replay particular import
+        # later to restore the data consistency.
         #
-        #   UsersIndex::User.import                          # imports default data set
-        #   UsersIndex::User.import User.active              # imports active users
-        #   UsersIndex::User.import [1, 2, 3]                # imports users with specified ids
-        #   UsersIndex::User.import users                    # imports users collection
-        #   UsersIndex::User.import suffix: Time.now.to_i    # imports data to index with specified suffix if such exists
-        #   UsersIndex::User.import refresh: false           # to disable index refreshing after import
-        #   UsersIndex::User.import journal: true            # import will record all the actions into special journal index
-        #   UsersIndex::User.import batch_size: 300          # import batch size
-        #   UsersIndex::User.import bulk_size: 10.megabytes  # import ElasticSearch bulk size in bytes
-        #   UsersIndex::User.import consistency: :quorum     # explicit write consistency setting for the operation (one, quorum, all)
-        #   UsersIndex::User.import replication: :async      # explicitly set the replication type (sync, async)
+        # Performs partial index update using `update` bulk action if any `fields` are
+        # specified. Note that if document doesn't exist yet, an error will be raised
+        # by ES, but import catches this an errors and performs full indexing
+        # for the corresponding documents. This feature can be disabled by setting
+        # `update_failover` to `false`.
         #
-        # See adapters documentation for more details.
+        # Utilizes `ActiveSupport::Notifications`, so it is possible to get imported
+        # objects later by listening to the `import_objects.chewy` queue. It is also
+        # possible to get the list of occured errors from the payload if something
+        # went wrong.
         #
+        # Import can also be run in parallel using the Parallel gem functionality.
+        #
+        # @example
+        #   UsersIndex::User.import(parallel: true) # imports everything in parallel with automatic workers number
+        #   UsersIndex::User.import(parallel: 3) # using 3 workers
+        #   UsersIndex::User.import(parallel: {in_threads: 10}) # in 10 threads
+        #
+        # @see https://github.com/elastic/elasticsearch-ruby/blob/master/elasticsearch-api/lib/elasticsearch/api/actions/bulk.rb
+        # @param collection [Array<Object>] and array or anything to import
+        # @param options [Hash{Symbol => Object}] besides specific import options, it accepts all the options suitable for the bulk API call like `refresh` or `timeout`
+        # @option options [String] suffix an index name suffix, used for zero-downtime reset mostly, no suffix by default
+        # @option options [Integer] bulk_size bulk API chunk size in bytes; if passed, the request is performed several times for each chunk, empty by default
+        # @option options [Integer] batch_size passed to the adapter import method, used to split imported objects in chunks, 1000 by default
+        # @option options [true, false] journal enables imported objects journaling, false by default
+        # @option options [Array<Symbol, String>] update_fields list of fields for the partial import, empty by default
+        # @option options [true, false] update_failover enables full objects reimport in cases of partial update errors, `true` by default
+        # @option options [true, Integer, Hash] parallel enables parallel import processing with the Parallel gem, accepts the number of workers or any Parallel gem acceptable options
+        # @return [true, false] false in case of errors
         def import(*args)
-          import_options = args.extract_options!
-          import_options.reverse_merge! _default_import_options
-          bulk_options = import_options.reject { |k, _| !BULK_OPTIONS.include?(k) }.reverse_merge!(refresh: true)
-
-          index.create!(bulk_options.slice(:suffix)) unless index.exists?
-
-          ActiveSupport::Notifications.instrument 'import_objects.chewy', type: self do |payload|
-            adapter.import(*args, import_options) do |action_objects|
-              journal = Chewy::Journal.new(self)
-              journal.add(action_objects) if import_options.fetch(:journal) { journal? }
-
-              indexed_objects = build_root.parent_id && fetch_indexed_objects(action_objects.values.flatten)
-              body = bulk_body(action_objects, indexed_objects)
-
-              errors = bulk(bulk_options.merge(body: body, journal: journal)) if body.present?
-
-              fill_payload_import payload, action_objects
-              fill_payload_errors payload, errors if errors.present?
-              !errors.present?
-            end
-          end
+          import_routine(*args).blank?
         end
 
-        # Perform import operation for specified documents.
-        # Raises Chewy::ImportFailed exception in case of import errors.
-        # Options are completely the same as for `import` method
-        # See adapters documentation for more details.
+        # @!method import!(*collection, **options)
+        # (see #import)
+        #
+        # The only difference from {#import} is that it raises an exception
+        # in case of any import errors.
         #
+        # @raise [Chewy::ImportFailed] in case of errors
         def import!(*args)
-          errors = nil
-          subscriber = ActiveSupport::Notifications.subscribe('import_objects.chewy') do |*notification_args|
-            errors = notification_args.last[:errors]
-          end
-          import(*args)
+          errors = import_routine(*args)
           raise Chewy::ImportFailed.new(self, errors) if errors.present?
           true
-        ensure
-          ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
         end
 
-        # Wraps elasticsearch-ruby client indices bulk method.
-        # Adds `:suffix` option to bulk import to index with specified suffix.
-        def bulk(options = {})
-          suffix = options.delete(:suffix)
-          bulk_size = options.delete(:bulk_size)
-          body = options.delete(:body)
-          journal = options.delete(:journal)
-          header = { index: index.build_index_name(suffix: suffix), type: type_name }
-
-          bodies = if bulk_size
-            bulk_size -= 1.kilobyte # 1 kilobyte for request header and newlines
-            raise ArgumentError, 'Import `:bulk_size` can\'t be less than 1 kilobyte' if bulk_size <= 0
-
-            entries = body.each_with_object(['']) do |entry, result|
-              operation, meta = entry.to_a.first
-              data = meta.delete(:data)
-              entry = [{ operation => meta }, data].compact.map(&:to_json).join("\n")
-
-              raise ArgumentError, 'Import `:bulk_size` seems to be less than entry size' if entry.bytesize > bulk_size
-
-              if result.last.bytesize + entry.bytesize > bulk_size
-                result.push(entry)
-              else
-                result[-1] = [result[-1], entry].delete_if(&:blank?).join("\n")
-              end
-            end
-            entries.map { |entry| entry + "\n" }
-          else
-            [body]
-          end
-
-          if journal.any_records?
-            Chewy::Journal.create
-            bodies += [journal.bulk_body]
-          end
-
-          items = bodies.map do |item_body|
-            result = client.bulk options.merge(header).merge(body: item_body)
-            result.try(:[], 'items') || []
-          end.flatten
+        # Wraps elasticsearch API bulk method, adds additional features like
+        # `bulk_size` and `suffix`.
+        #
+        # @see https://github.com/elastic/elasticsearch-ruby/blob/master/elasticsearch-api/lib/elasticsearch/api/actions/bulk.rb
+        # @see Chewy::Type::Import::Bulk
+        # @param options [Hash{Symbol => Object}] besides specific import options, it accepts all the options suitable for the bulk API call like `refresh` or `timeout`
+        # @option options [String] suffix bulk API chunk size in bytes; if passed, the request is performed several times for each chunk, empty by default
+        # @option options [Integer] bulk_size bulk API chunk size in bytes; if passed, the request is performed several times for each chunk, empty by default
+        # @option options [Array<Hash>] body elasticsearch API bulk method body
+        # @return [Hash] tricky transposed errors hash, empty if everything is fine
+        def bulk(**options)
+          error_items = BulkRequest.new(self, **options).perform(options[:body])
           Chewy.wait_for_status
 
-          extract_errors items
+          payload_errors(error_items)
         end
 
-        def journal?
-          _default_import_options.fetch(:journal) { Chewy.configuration[:journal] }
+        # Composes a single document from the passed object. Uses either witchcraft
+        # or normal composing under the hood.
+        #
+        # @param object [Object] a data source object
+        # @param crutches [Object] optional crutches object; if ommited - a crutch for the single passed object is created as a fallback
+        # @param fields [Array<Symbol>] and array of fields to restrict the generated document
+        # @return [Hash] a JSON-ready hash
+        def compose(object, crutches = nil, fields: [])
+          crutches ||= Chewy::Type::Crutch::Crutches.new self, [object]
+
+          if witchcraft? && build_root.children.present?
+            cauldron(fields: fields).brew(object, crutches)
+          else
+            build_root.compose(object, crutches, fields: fields)
+          end
         end
 
       private
 
-        def bulk_body(action_objects, indexed_objects = nil)
-          action_objects.flat_map do |action, objects|
-            method = "#{action}_bulk_entry"
-            crutches = Chewy::Type::Crutch::Crutches.new self, objects
-            objects.flat_map { |object| send(method, object, indexed_objects, crutches) }
-          end
-        end
-
-        def delete_bulk_entry(object, indexed_objects = nil, _crutches = nil)
-          entry = {}
+        def import_routine(*args)
+          routine = Routine.new(self, args.extract_options!)
+          routine.create_indexes!
 
-          if root_object.id
-            entry[:_id] = root_object.compose_id(object)
+          if routine.parallel_options
+            import_parallel(args, routine)
           else
-            entry[:_id] = object.id if object.respond_to?(:id)
-            entry[:_id] ||= object[:id] || object['id'] if object.is_a?(Hash)
-            entry[:_id] ||= object
-            entry[:_id] = entry[:_id].to_s if defined?(BSON) && entry[:_id].is_a?(BSON::ObjectId)
+            import_linear(args, routine)
           end
-
-          if root_object.parent_id
-            existing_object = entry[:_id].present? && indexed_objects && indexed_objects[entry[:_id].to_s]
-            return [] unless existing_object
-            entry[:parent] = existing_object[:parent]
-          end
-
-          [{ delete: entry }]
         end
 
-        def index_bulk_entry(object, indexed_objects = nil, crutches = nil)
-          entry = {}
-
-          if root_object.id
-            entry[:_id] = root_object.compose_id(object)
-          else
-            entry[:_id] = object.id if object.respond_to?(:id)
-            entry[:_id] ||= object[:id] || object['id'] if object.is_a?(Hash)
-            entry[:_id] = entry[:_id].to_s if defined?(BSON) && entry[:_id].is_a?(BSON::ObjectId)
-          end
-          entry.delete(:_id) if entry[:_id].blank?
-
-          if root_object.parent_id
-            entry[:parent] = root_object.compose_parent(object)
-            existing_object = entry[:_id].present? && indexed_objects && indexed_objects[entry[:_id].to_s]
+        def import_linear(objects, routine)
+          ActiveSupport::Notifications.instrument 'import_objects.chewy', type: self do |payload|
+            adapter.import(*objects, routine.options) do |action_objects|
+              routine.process(**action_objects)
+            end
+            routine.perform_bulk(routine.leftovers)
+            payload[:import] = routine.stats
+            payload[:errors] = payload_errors(routine.errors) if routine.errors.present?
+            payload[:errors]
           end
+        end
 
-          entry[:data] = object_data(object, crutches)
+        def import_parallel(objects, routine)
+          raise "The `parallel` gem is required for parallel import, please add `gem 'parallel'` to your Gemfile" unless '::Parallel'.safe_constantize
 
-          if existing_object && entry[:parent].to_s != existing_object[:parent]
-            [{ delete: entry.except(:data).merge(parent: existing_object[:parent]) }, { index: entry }]
-          else
-            [{ index: entry }]
-          end
-        end
+          ActiveSupport::Notifications.instrument 'import_objects.chewy', type: self do |payload|
+            batches = adapter.import_references(*objects, routine.options.slice(:batch_size)).to_a
 
-        def fill_payload_import(payload, action_objects)
-          imported = Hash[action_objects.map { |action, objects| [action, objects.count] }]
-          imported.each do |action, count|
-            payload[:import] ||= {}
-            payload[:import][action] ||= 0
-            payload[:import][action] += count
-          end
-        end
+            ::ActiveRecord::Base.connection.close if defined?(::ActiveRecord::Base)
+            results = ::Parallel.map(batches, routine.parallel_options, &IMPORT_WORKER.curry[self, routine.options])
+            ::ActiveRecord::Base.connection.reconnect! if defined?(::ActiveRecord::Base)
+            errors, import, leftovers = process_parallel_import_results(results)
 
-        def fill_payload_errors(payload, import_errors)
-          import_errors.each do |action, action_errors|
-            action_errors.each do |error, documents|
-              payload[:errors] ||= {}
-              payload[:errors][action] ||= {}
-              payload[:errors][action][error] ||= []
-              payload[:errors][action][error] |= documents
+            if leftovers.present?
+              batches = leftovers.each_slice(routine.options[:batch_size])
+              results = ::Parallel.map(batches, routine.parallel_options, &LEFTOVERS_WORKER.curry[self, routine.options])
+              errors.concat(results.flatten(1))
             end
-          end
-        end
 
-        def object_data(object, crutches = nil)
-          if witchcraft?
-            cauldron.brew(object, crutches)
-          else
-            build_root.compose(object, crutches)[type_name.to_s]
+            payload[:import] = import
+            payload[:errors] = payload_errors(errors) if errors.present?
+            payload[:errors]
           end
         end
 
-        def extract_errors(items)
-          items = items.each.with_object({}) do |item, memo|
-            action = item.keys.first.to_sym
-            data = item.values.first
-            if data['error']
-              (memo[action] ||= []).push(action: action, id: data['_id'], error: data['error'])
-            end
+        def process_parallel_import_results(results)
+          results.each_with_object([[], {}, []]) do |r, (e, i, l)|
+            e.concat(r[:errors])
+            i.merge!(r[:import]) { |_k, v1, v2| v1.to_i + v2.to_i }
+            l.concat(r[:leftovers])
           end
-
-          items.map do |action, action_items|
-            errors = action_items.group_by { |item| item[:error] }.map do |error, error_items|
-              { error => error_items.map { |item| item[:id] } }
-            end.reduce(&:merge)
-            { action => errors }
-          end.reduce(&:merge) || {}
         end
 
-        def fetch_indexed_objects(objects)
-          ids = objects.map { |object| object.respond_to?(:id) ? object.id : object }
-          result = client.search index: index_name,
-                                 type: type_name,
-                                 fields: '_parent',
-                                 body: { filter: { ids: { values: ids } } },
-                                 search_type: 'scan',
-                                 scroll: '1m'
-
-          indexed_objects = {}
+        def payload_errors(errors)
+          errors.each_with_object({}) do |error, result|
+            action = error.keys.first.to_sym
+            item = error.values.first
+            error = item['error']
+            id = item['_id']
 
-          while (result = client.scroll(scroll_id: result['_scroll_id'], scroll: '1m'))
-            break if result['hits']['hits'].empty?
-
-            result['hits']['hits'].map do |hit|
-              parent = hit.key?('_parent') ? hit['_parent'] : hit['fields']['_parent']
-              indexed_objects[hit['_id']] = { parent: parent }
-            end
+            result[action] ||= {}
+            result[action][error] ||= []
+            result[action][error].push(id)
           end
-
-          indexed_objects
         end
       end
     end