rocketjob 5.4.1 → 6.0.0

data/lib/rocket_job/batch/io.rb

@@ -9,34 +9,66 @@ module RocketJob
       # Returns [RocketJob::Sliced::Input] input collection for holding input slices
       #
       # Parameters:
-      #   category [Symbol]
-      #     The name of the category to access or upload data into
+      #   category [Symbol|RocketJob::Category::Input]
+      #     The category or the name of the category to access or upload data into
       #     Default: None ( Uses the single default input collection for this job )
       #     Validates: This value must be one of those listed in #input_categories
       def input(category = :main)
-        unless input_categories.include?(category) || (category == :main)
-          raise "Category #{category.inspect}, must be registered in input_categories: #{input_categories.inspect}"
-        end
+        category = input_category(category)
 
-        (@inputs ||= {})[category] ||= RocketJob::Sliced.factory(:input, category, self)
+        (@inputs ||= {})[category.name] ||= category.data_store(self)
       end
 
       # Returns [RocketJob::Sliced::Output] output collection for holding output slices
       # Returns nil if no output is being collected
       #
       # Parameters:
-      #   category [Symbol]
-      #     The name of the category to access or download data from
+      #   category [Symbol|RocketJob::Category::Input]
+      #     The category or the name of the category to access or download data from
       #     Default: None ( Uses the single default output collection for this job )
       #     Validates: This value must be one of those listed in #output_categories
       def output(category = :main)
-        unless output_categories.include?(category) || (category == :main)
-          raise "Category #{category.inspect}, must be registered in output_categories: #{output_categories.inspect}"
-        end
+        category = output_category(category)
 
-        (@outputs ||= {})[category] ||= RocketJob::Sliced.factory(:output, category, self)
+        (@outputs ||= {})[category.name] ||= category.data_store(self)
       end
 
+      # Rapidly upload individual records in batches.
+      #
+      # Operates directly on a Mongo Collection to avoid the overhead of creating Mongoid objects
+      # for each and every row.
+      #
+      #   input_category(:my_lookup).find(id: 123).first
+      #
+      # Lookup collection.
+      #
+      # Upload side / secondary lookup tables that can be accessed during job processing.
+      #
+      # Example:
+      #   lookup_collection(:my_lookup).upload do |io|
+      #     io << {id: 123, data: "first record"}
+      #     io << {id: 124, data: "second record"}
+      #   end
+      #
+      # Parameters:
+      #   category [Symbol|RocketJob::Category::Input]
+      #     The category or the name of the category to access or download data from
+      #     Default: None ( Uses the single default output collection for this job )
+      #     Validates: This value must be one of those listed in #input_categories
+      # def lookup_collection(category = :main)
+      #   category = input_category(category) unless category.is_a?(Category::Input)
+      #
+      #   collection = (@lookup_collections ||= {})[category.name]
+      #
+      #   unless collection
+      #     collection_name = "rocket_job.inputs.#{id}"
+      #     collection_name << ".#{category.name}" unless category.name == :main
+      #
+      #     @lookup_collections[category.name] ||=
+      #       LookupCollection.new(Sliced::Slice.collection.database, collection_name)
+      #   end
+      # end
+
       # Upload the supplied file, io, IOStreams::Path, or IOStreams::Stream.
       #
       # Returns [Integer] the number of records uploaded.
@@ -65,6 +97,11 @@ module RocketJob
       #       Parses each line from the file into a Hash and uploads each hash for processing by workers.
       #     See IOStreams::Stream#each.
       #
+      #   category [Symbol|RocketJob::Category::Input]
+      #     The category or the name of the category to access or download data from
+      #     Default: None ( Uses the single default output collection for this job )
+      #     Validates: This value must be one of those listed in #input_categories
+      #
       # Example:
       #   # Load plain text records from a file
       #   job.upload('hello.csv')
@@ -113,29 +150,7 @@ module RocketJob
       # * If an io stream is supplied, it is read until it returns nil.
       # * Only use this method for UTF-8 data, for binary data use #input_slice or #input_records.
       # * CSV parsing is slow, so it is usually left for the workers to do.
-      def upload(stream = nil, file_name: nil, category: :main, stream_mode: :line, on_first: nil, **args, &block)
-        raise(ArgumentError, "Either stream, or a block must be supplied") unless stream || block
-
-        stream_mode = stream_mode.to_sym
-        # Backward compatibility with existing v4 jobs
-        stream_mode = :array if stream_mode == :row
-        stream_mode = :hash if stream_mode == :record
-
-        count             =
-          if block
-            input(category).upload(on_first: on_first, &block)
-          else
-            path                  = IOStreams.new(stream)
-            path.file_name        = file_name if file_name
-            self.upload_file_name = path.file_name
-            input(category).upload(on_first: on_first) do |io|
-              path.each(stream_mode, **args) { |line| io << line }
-            end
-          end
-        self.record_count = (record_count || 0) + count
-        count
-      end
-
+      #
       # Upload results from an Arel into RocketJob::SlicedJob.
       #
       # Params
@@ -144,6 +159,9 @@ module RocketJob
       #     and uploaded into the job
       #     These columns are automatically added to the select list to reduce overhead
       #
+      #   category [Symbol|RocketJob::Category::Input]
+      #     The category or the name of the category to upload to.
+      #
       # If a Block is supplied it is passed the model returned from the database and should
       # return the work item to be uploaded into the job.
       #
@@ -159,18 +177,13 @@ module RocketJob
       #
       # Example: Upload user_name and zip_code
       #   arel = User.where(country_code: 'US')
-      #   job.upload_arel(arel, :user_name, :zip_code)
+      #   job.upload_arel(arel, columns: [:user_name, :zip_code])
       #
       # Notes:
       # * Only call from one thread at a time against a single instance of this job.
       # * The record_count for the job is set to the number of records returned by the arel.
       # * If an exception is raised while uploading data, the input collection is cleared out
       #   so that if a job is retried during an upload failure, data is not duplicated.
-      def upload_arel(arel, *column_names, category: :main, &block)
-        count             = input(category).upload_arel(arel, *column_names, &block)
-        self.record_count = (record_count || 0) + count
-        count
-      end
 
       # Upload the result of a MongoDB query to the input collection for processing
       # Useful when an entire MongoDB collection, or part thereof needs to be
@@ -198,30 +211,25 @@ module RocketJob
       #   criteria = User.where(state: 'FL')
       #   job.record_count = job.upload_mongo_query(criteria)
       #
-      # Example: Upload just the supplied column
+      # Example: Upload only the specified column(s)
       #   criteria = User.where(state: 'FL')
-      #   job.record_count = job.upload_mongo_query(criteria, :zip_code)
+      #   job.record_count = job.upload_mongo_query(criteria, columns: [:zip_code])
       #
       # Notes:
       # * Only call from one thread at a time against a single instance of this job.
       # * The record_count for the job is set to the number of records returned by the monqo query.
       # * If an exception is raised while uploading data, the input collection is cleared out
       #   so that if a job is retried during an upload failure, data is not duplicated.
-      def upload_mongo_query(criteria, *column_names, category: :main, &block)
-        count             = input(category).upload_mongo_query(criteria, *column_names, &block)
-        self.record_count = (record_count || 0) + count
-        count
-      end
 
       # Upload sliced range of integer requests as arrays of start and end ids.
       #
-      # Returns [Integer] last_id - start_id + 1.
+      # Returns [Integer] the number of slices uploaded.
       #
       # Uploads one range per slice so that the response can return multiple records
       # for each slice processed
       #
       # Example
-      #   job.slice_size = 100
+      #   job.input_category.slice_size = 100
       #   job.upload_integer_range(200, 421)
       #
       #   # Equivalent to calling:
@@ -234,17 +242,11 @@ module RocketJob
       # * The record_count for the job is set to: last_id - start_id + 1.
       # * If an exception is raised while uploading data, the input collection is cleared out
       #   so that if a job is retried during an upload failure, data is not duplicated.
-      def upload_integer_range(start_id, last_id, category: :main)
-        input(category).upload_integer_range(start_id, last_id)
-        count             = last_id - start_id + 1
-        self.record_count = (record_count || 0) + count
-        count
-      end
 
       # Upload sliced range of integer requests as an arrays of start and end ids
       # starting with the last range first
       #
-      # Returns [Integer] last_id - start_id + 1.
+      # Returns [Integer] the number of slices uploaded.
       #
       # Uploads one range per slice so that the response can return multiple records
       # for each slice processed.
@@ -253,7 +255,7 @@ module RocketJob
       # in a database based on the id column
       #
       # Example
-      #   job.slice_size = 100
+      #   job.input_category.slice_size = 100
       #   job.upload_integer_range_in_reverse_order(200, 421)
       #
       #   # Equivalent to calling:
@@ -266,14 +268,102 @@ module RocketJob
       # * The record_count for the job is set to: last_id - start_id + 1.
       # * If an exception is raised while uploading data, the input collection is cleared out
       #   so that if a job is retried during an upload failure, data is not duplicated.
-      def upload_integer_range_in_reverse_order(start_id, last_id, category: :main)
-        input(category).upload_integer_range_in_reverse_order(start_id, last_id)
-        count             = last_id - start_id + 1
+
+      def upload(object = nil, category: :main, file_name: nil, stream_mode: nil, on_first: nil, columns: nil, slice_batch_size: nil, **args, &block)
+        input_collection = input(category)
+
+        if block
+          raise(ArgumentError, "Cannot supply both an object to upload, and a block.") if object
+          if stream_mode || columns || slice_batch_size || args.size > 0
+            raise(ArgumentError, "Unknown keyword arguments when uploading a block. Only accepts :category, :file_name, or :on_first")
+          end
+
+          category           = input_category(category)
+          category.file_name = file_name if file_name
+
+          # Extract the header line during the upload when applicable.
+          extract_header = category.extract_header_callback(on_first)
+
+          count             = input_collection.upload(on_first: extract_header, slice_batch_size: slice_batch_size, &block)
+          self.record_count = (record_count || 0) + count
+          return count
+        end
+
+        count =
+          case object
+          when Range
+            if file_name || stream_mode || on_first || args.size > 0
+              raise(ArgumentError, "Unknown keyword arguments when uploading a Range. Only accepts :category, :columns, or :slice_batch_size")
+            end
+
+            first = object.first
+            last  = object.last
+            if first < last
+              input_collection.upload_integer_range(first, last, slice_batch_size: slice_batch_size || 1_000)
+            else
+              input_collection.upload_integer_range_in_reverse_order(last, first, slice_batch_size: slice_batch_size || 1_000)
+            end
+          when Mongoid::Criteria
+            if file_name || stream_mode || on_first || args.size > 0
+              raise(ArgumentError, "Unknown keyword arguments when uploading a Mongoid::Criteria. Only accepts :category, :columns, or :slice_batch_size")
+            end
+
+            input_collection.upload_mongo_query(object, columns: columns, slice_batch_size: slice_batch_size, &block)
+          when defined?(ActiveRecord::Relation) ? ActiveRecord::Relation : false
+            if file_name || stream_mode || on_first || args.size > 0
+              raise(ArgumentError, "Unknown keyword arguments when uploading an ActiveRecord::Relation. Only accepts :category, :columns, or :slice_batch_size")
+            end
+
+            input_collection.upload_arel(object, columns: columns, slice_batch_size: slice_batch_size, &block)
+
+          else
+            raise(ArgumentError, "Unknown keyword argument :columns when uploading a file") if columns
+
+            category = input_category(category)
+
+            # Extract the header line during the upload when applicable.
+            extract_header = category.extract_header_callback(on_first)
+            path = category.upload_path(object, original_file_name: file_name)
+
+            input_collection.upload(on_first: extract_header, slice_batch_size: slice_batch_size) do |io|
+              path.each(stream_mode || :line, **args) { |line| io << line }
+            end
+
+          end
+
+        self.record_count = (record_count || 0) + count
+        count
+      end
+
+      # @deprecated
+      def upload_arel(arel, *column_names, category: :main, &block)
+        count             = input(category).upload_arel(arel, columns: column_names, &block)
+        self.record_count = (record_count || 0) + count
+        count
+      end
+
+      # @deprecated
+      def upload_mongo_query(criteria, *column_names, category: :main, &block)
+        count             = input(category).upload_mongo_query(criteria, columns: column_names, &block)
         self.record_count = (record_count || 0) + count
         count
       end
 
-      # Upload the supplied slices for processing by workers
+      # @deprecated
+      def upload_integer_range(start_id, last_id, category: :main, slice_batch_size: 1_000)
+        count             = input(category).upload_integer_range(start_id, last_id, slice_batch_size: slice_batch_size)
+        self.record_count = (record_count || 0) + count
+        count
+      end
+
+      # @deprecated
+      def upload_integer_range_in_reverse_order(start_id, last_id, category: :main, slice_batch_size: 1_000)
+        count             = input(category).upload_integer_range_in_reverse_order(start_id, last_id, slice_batch_size: slice_batch_size)
+        self.record_count = (record_count || 0) + count
+        count
+      end
+
+      # Upload the supplied slice for processing by workers
       #
       # Updates the record_count after adding the records
       #
@@ -285,12 +375,12 @@ module RocketJob
       #     For example the following types are not supported: Date
       #
       # Note:
-      #   The caller should honor `:slice_size`, the entire slice is loaded as-is.
+      #   The caller should implement `:slice_size`, since the entire slice is saved as-is.
       #
       # Note:
       #   Not thread-safe. Only call from one thread at a time
-      def upload_slice(slice)
-        input.insert(slice)
+      def upload_slice(slice, category: :main)
+        input(category).insert(slice)
         count             = slice.size
         self.record_count = (record_count || 0) + count
         count
@@ -353,56 +443,34 @@ module RocketJob
       def download(stream = nil, category: :main, header_line: nil, **args, &block)
         raise "Cannot download incomplete job: #{id}. Currently in state: #{state}-#{sub_state}" if rocket_job_processing?
 
-        return output(category).download(header_line: header_line, &block) if block
+        category           = output_category(category) unless category.is_a?(Category::Output)
+        output_collection  = output(category)
+
+        # Store the output file name in the category
+        category.file_name = stream if !block && (stream.is_a?(String) || stream.is_a?(IOStreams::Path))
 
-        output_collection = output(category)
+        header_line ||= category.render_header
 
-        if output_collection.binary?
-          IOStreams.new(stream).stream(:none).writer(**args) do |io|
-            raise(ArgumenError, "A `header_line` is not supported with binary output collections") if header_line
+        return output_collection.download(header_line: header_line, &block) if block
 
-            output_collection.download { |record| io << record[:binary] }
+        raise(ArgumentError, "Missing mandatory `stream` or `category.file_name`") unless stream || category.file_name
+
+        if output_collection.slice_class.binary_format
+          binary_header_line = output_collection.slice_class.to_binary(header_line) if header_line
+
+          # Don't overwrite supplied stream options if any
+          stream = stream&.is_a?(IOStreams::Stream) ? stream.dup : IOStreams.new(category.file_name)
+          stream.remove_from_pipeline(output_collection.slice_class.binary_format)
+          stream.writer(**args) do |io|
+            # TODO: Binary formats should return the record count, instead of the slice count.
+            output_collection.download(header_line: binary_header_line) { |record| io.write(record) }
           end
         else
-          IOStreams.new(stream).writer(:line, **args) do |io|
+          IOStreams.new(stream || category.file_name).writer(:line, **args) do |io|
             output_collection.download(header_line: header_line) { |record| io << record }
           end
         end
       end
-
-      # Writes the supplied result, Batch::Result or Batch::Results to the relevant collections.
-      #
-      # If a block is supplied, the block is supplied with a writer that should be used to
-      # accumulate the results.
-      #
-      # Examples
-      #
-      # job.write_output('hello world')
-      #
-      # job.write_output do |writer|
-      #   writer << 'hello world'
-      # end
-      #
-      # job.write_output do |writer|
-      #   result = RocketJob::Batch::Results
-      #   result << RocketJob::Batch::Result.new(:main, 'hello world')
-      #   result << RocketJob::Batch::Result.new(:errors, 'errors')
-      #   writer << result
-      # end
-      #
-      # result = RocketJob::Batch::Results
-      # result << RocketJob::Batch::Result.new(:main, 'hello world')
-      # result << RocketJob::Batch::Result.new(:errors, 'errors')
-      # job.write_output(result)
-      def write_output(result = nil, input_slice = nil, &block)
-        if block
-          RocketJob::Sliced::Writer::Output.collect(self, input_slice, &block)
-        else
-          raise(ArgumentError, "result parameter is required when no block is supplied") unless result
-
-          RocketJob::Sliced::Writer::Output.collect(self, input_slice) { |writer| writer << result }
-        end
-      end
     end
   end
 end

data/lib/rocket_job/batch/model.rb

@@ -11,46 +11,6 @@ module RocketJob
         #
         # The following attributes are set when the job is created
 
-        # Number of records to include in each slice that is processed
-        # Note:
-        #   slice_size is only used by SlicedJob#upload & Sliced::Input#upload
-        #   When slices are supplied directly, their size is not modified to match this number
-        field :slice_size, type: Integer, default: 100, class_attribute: true, user_editable: true, copy_on_restart: true
-
-        # Whether to retain nil results.
-        #
-        # Only applicable if `collect_output` is `true`
-        # Set to `false` to prevent collecting output from the perform
-        # method when it returns `nil`.
-        field :collect_nil_output, type: Boolean, default: true, class_attribute: true
-
-        # Optional Array<Symbol> list of categories that this job can output to
-        #
-        # By using categories the output from #perform can be placed in different
-        # output collections, and therefore different output files
-        #
-        # Categories must be declared in advance to avoid a #perform method
-        # accidentally writing its results to an unknown category
-        field :output_categories, type: Array, default: [:main], class_attribute: true
-
-        # Optional Array<Symbol> list of categories that this job can load input data into
-        field :input_categories, type: Array, default: [:main], class_attribute: true
-
-        # The file name of the uploaded file, if any.
-        # Set by #upload if a file name was supplied, but can also be set explicitly.
-        # May or may not include the fully qualified path name.
-        field :upload_file_name, type: String
-
-        # Compress uploaded records.
-        # The fields are not affected in any way, only the data stored in the
-        # records and results collections will compressed
-        field :compress, type: Object, default: false, class_attribute: true
-
-        # Encrypt uploaded records.
-        # The fields are not affected in any way, only the data stored in the
-        # records and results collections will be encrypted
-        field :encrypt, type: Object, default: false, class_attribute: true
-
         #
         # Values that jobs can also update during processing
         #
@@ -69,30 +29,7 @@ module RocketJob
 
         # Breaks the :running state up into multiple sub-states:
         #   :running -> :before -> :processing -> :after -> :complete
-        field :sub_state, type: Symbol
-
-        validates_presence_of :slice_size
-
-        validates_each :output_categories, :input_categories do |record, attr, value|
-          # Under some circumstances ActiveModel is passing in a nil value even though the
-          # attributes have default values
-          Array(value).each do |category|
-            record.errors.add(attr, "must only contain Symbol values") unless category.is_a?(Symbol)
-            unless category.to_s =~ /\A[a-z_0-9]+\Z/
-              record.errors.add(attr, "must only consist of lowercase characters, digits, and _")
-            end
-          end
-        end
-      end
-
-      # Returns [true|false] whether the slices for this job are encrypted
-      def encrypted?
-        encrypt == true
-      end
-
-      # Returns [true|false] whether the slices for this job are compressed
-      def compressed?
-        compress == true
+        field :sub_state, type: Mongoid::StringifiedSymbol
       end
 
       # Returns [Integer] percent of records completed so far
@@ -102,10 +39,10 @@ module RocketJob
         return 0 unless record_count.to_i.positive?
 
         # Approximate number of input records
-        input_records = input.count.to_f * slice_size
+        input_records = input.count.to_f * input_category.slice_size
         if input_records > record_count
           # Sanity check in case slice_size is not being adhered to
-          99
+          0
         else
           ((1.0 - (input_records.to_f / record_count)) * 100).to_i
         end
@@ -120,6 +57,10 @@ module RocketJob
           h["active_slices"] = worker_count
           h["failed_slices"] = input.failed.count
           h["queued_slices"] = input.queued.count
+          output_categories.each do |category|
+            name_str                      = category.name == :main ? "" : "_#{category.name}"
+            h["output_slices#{name_str}"] = output(category).count
+          end
           # Very high level estimated time left
           if record_count && running? && record_count.positive?
             percent = percent_complete
@@ -129,10 +70,9 @@ module RocketJob
             end
           end
         elsif completed?
-          secs = seconds.to_f
+          secs                  = seconds.to_f
           h["records_per_hour"] = ((record_count.to_f / secs) * 60 * 60).round if record_count&.positive? && (secs > 0.0)
         end
-        h["output_slices"] = output.count if collect_output? && !completed?
         h.merge!(super(time_zone))
         h.delete("result")
         # Worker name should be retrieved from the slices when processing
@@ -172,6 +112,18 @@ module RocketJob
         @worker_count_last = Time.now.to_i
         @worker_count
       end
+
+      # @deprecated
+      # For backward compatibility
+      def upload_file_name
+        input_category.file_name
+      end
+
+      # @deprecated
+      # For backward compatibility
+      def upload_file_name=(upload_file_name)
+        input_category.file_name = upload_file_name
+      end
     end
   end
 end