rocketjob 3.5.2 → 4.0.0

data/lib/rocket_job/sliced/output.rb

@@ -0,0 +1,99 @@
+require 'tempfile'
+
+module RocketJob
+  module Sliced
+    class Output < Slices
+      # Write this output collection to the specified file/io stream
+      #
+      # Returns [Integer] the number of records returned from the collection
+      #
+      # Parameters
+      #   file_name_or_io [String|IO]
+      #     The file_name of the file to write to, or an IO Stream that implements
+      #     #write.
+      #
+      #   options:
+      #     streams [Symbol|Array]
+      #       The formats/streams that be used to convert the data whilst it is
+      #       being written.
+      #       When nil, `file_name_or_io` will be inspected to try and determine what
+      #       streams should be applied.
+      #       Default: nil
+      #
+      #     Any other option that can be supplied to IOStreams::Line::Writer
+      #
+      # Stream types / extensions supported:
+      #   .zip       Zip File                                   [ :zip ]
+      #   .gz, .gzip GZip File                                  [ :gzip ]
+      #   .enc       File Encrypted using symmetric encryption  [ :enc ]
+      #
+      # When a file is encrypted, it may also be compressed:
+      #   .zip.enc  [ :zip, :enc ]
+      #   .gz.enc   [ :gz,  :enc ]
+      #
+      # Example: Zip
+      #   # Since csv is not known to RocketJob it is ignored
+      #   job.output.download('myfile.csv.zip')
+      #
+      # Example: Encrypted Zip
+      #   job.output.download('myfile.csv.zip.enc')
+      #
+      # Example: Explicitly set the streams
+      #   job.output.download('myfile.ze', streams: [:zip, :enc])
+      #
+      # Example: Supply custom options
+      #   job.output.download('myfile.csv.enc', streams: [enc: { compress: true }])
+      #
+      # Example: Supply custom options
+      #   job.output.download('myfile.csv.zip', streams: [ zip: { zip_file_name: 'myfile.csv' } ])
+      #
+      # Example: Extract streams from filename but write to a temp file
+      #   t = Tempfile.new('my_project')
+      #   job.output.download(t.to_path, file_name: 'myfile.gz.enc')
+      #
+      # Example: Add a header and/or trailer record to the downloaded file:
+      #   IOStreams.writer('/tmp/file.txt.gz') do |writer|
+      #     writer << "Header\n"
+      #     job.download do |line|
+      #       writer << line
+      #     end
+      #     writer << "Trailer\n"
+      #   end
+      #
+      # Notes:
+      # - The records are returned in '_id' order. Usually this is the order in
+      #   which the records were originally loaded.
+      def download(file_name_or_io = nil, header_line: nil, **args)
+        raise(ArgumentError, 'Either file_name_or_io, or a block must be supplied') unless file_name_or_io || block_given?
+
+        record_count = 0
+
+        if block_given?
+          # Write the header line
+          yield(header_line) if header_line
+
+          # Call the supplied block for every record returned
+          each do |slice|
+            slice.each do |record|
+              record_count += 1
+              yield(record)
+            end
+          end
+        else
+          IOStreams.line_writer(file_name_or_io, **args) do |io|
+            # Write the header line
+            io << header_line if header_line
+
+            each do |slice|
+              slice.each do |record|
+                record_count += 1
+                io << record
+              end
+            end
+          end
+        end
+        record_count
+      end
+    end
+  end
+end

data/lib/rocket_job/sliced/slice.rb

@@ -0,0 +1,166 @@
+require 'forwardable'
+module RocketJob
+  module Sliced
+    # A slice is an Array of Records, along with meta-data that is used
+    # or set during processing of the individual records
+    #
+    # Note: Do _not_ create instances of this model directly, go via Slice#new
+    #       so that the correct collection name is used.
+    #
+    # Example:
+    #   slice = RocketJob::Sliced::Slice.new
+    #   slice << 'first'
+    #   slice << 'second'
+    #   second = slice.at(1)
+    #
+    #   # The [] operator is for retrieving attributes:
+    #   slice['state']
+    #
+    class Slice
+      include RocketJob::Plugins::Document
+      include RocketJob::Plugins::StateMachine
+      extend Forwardable
+
+      store_in client: 'rocketjob_slices'
+
+      # The record number of the first record in this slice.
+      #
+      # Optional: If present the record_number is set while the job
+      #           is being processed.
+      field :first_record_number, type: Integer
+
+      #
+      # Read-only attributes
+      #
+
+      # Current state, as set by AASM
+      field :state, type: Symbol, default: :queued
+
+      # When processing started on this slice
+      field :started_at, type: Time
+
+      # Number of times that this job has failed to process
+      field :failure_count, type: Integer
+
+      # This name of the worker that this job is being processed by, or was processed by
+      field :worker_name, type: String
+
+      # The last exception for this slice if any
+      embeds_one :exception, class_name: 'RocketJob::JobException'
+
+      after_find :parse_records
+
+      # State Machine events and transitions
+      #
+      # Each slice is processed separately:
+      #   :queued -> :running -> :completed
+      #                       -> :failed     -> :running  ( manual )
+      #
+      # Slices are processed by ascending _id sort order
+      #
+      # Note:
+      #   Currently all slices are destroyed on completion, so no slices
+      #   are available in the completed state
+      aasm column: :state, whiny_persistence: true do
+        # Job has been created and is queued for processing ( Initial state )
+        state :queued, initial: true
+
+        # Job is running
+        state :running
+
+        # Job has completed processing ( End state )
+        state :completed
+
+        # Job failed to process and needs to be manually re-tried or aborted
+        state :failed
+
+        event :start, before: :set_started_at do
+          transitions from: :queued, to: :running
+        end
+
+        event :complete do
+          transitions from: :running, to: :completed
+        end
+
+        event :fail, before: :set_exception do
+          transitions from: :running, to: :failed
+          transitions from: :queued, to: :failed
+        end
+
+        event :retry do
+          transitions from: :failed, to: :queued
+        end
+      end
+
+      # `records` array has special handling so that it can be modified in place instead of having
+      # to replace the entire array every time. For example, when appending lines with `<<`.
+      def records
+        @records ||= []
+      end
+
+      # Replace the records within this slice
+      def records=(records)
+        raise(ArgumentError, "Cannot assign type: #{records.class.name} to records") unless records.is_a?(Array)
+
+        @records = records
+      end
+
+      def_instance_delegators :records, :each, :<<, :size, :concat, :at
+      def_instance_delegators :records, *(Enumerable.instance_methods - Module.methods)
+
+      # Fail this slice, along with the exception that caused the failure
+      def set_exception(exc = nil, record_number = nil)
+        if exc
+          self.exception          = JobException.from_exception(exc)
+          exception.worker_name   = worker_name
+          exception.record_number = record_number
+        end
+        self.failure_count = failure_count.to_i + 1
+        self.worker_name   = nil
+      end
+
+      # Returns [Hash] the slice as a Hash for storage purposes
+      # Compresses / Encrypts the slice according to the job setting
+      if ::Mongoid::VERSION.to_i >= 6
+        def as_attributes
+          attrs            = super
+          attrs['records'] = serialize_records if @records
+          attrs
+        end
+      else
+        def as_document
+          attrs            = super
+          attrs['records'] = serialize_records if @records
+          attrs
+        end
+      end
+
+      def inspect
+        "#{super[0...-1]}, records: #{@records.inspect}, collection_name: #{collection_name.inspect}>"
+      end
+
+      private
+
+      # Always add records to any updates.
+      def atomic_updates(*args)
+        r = super(*args)
+        if @records
+          (r['$set'] ||= {})['records'] = serialize_records
+        end
+        r
+      end
+
+      def parse_records
+        @records = attributes.delete('records')
+      end
+
+      def serialize_records
+        records.mongoize
+      end
+
+      def set_started_at
+        self.started_at = Time.now
+      end
+    end
+  end
+end

data/lib/rocket_job/sliced/slices.rb

@@ -0,0 +1,166 @@
+module RocketJob
+  module Sliced
+    class Slices
+      extend Forwardable
+      include Enumerable
+      include SemanticLogger::Loggable
+
+      attr_accessor :slice_class, :slice_size, :collection_name
+      attr_reader :all
+
+      # Parameters
+      #   name: [String]
+      #     Name of the collection to create
+      #   slice_size: [Integer]
+      #     Number of records to store in each slice
+      #     Default: 100
+      #   slice_class: [class]
+      #     Slice class to use to hold records.
+      #     Default: RocketJob::Sliced::Slice
+      def initialize(collection_name:, slice_class: Sliced::Slice, slice_size: 100)
+        @slice_class     = slice_class
+        @slice_size      = slice_size
+        @collection_name = collection_name
+        @all             = slice_class.with_collection(collection_name)
+      end
+
+      def new(params = {})
+        slice_class.new(params.merge(collection_name: collection_name))
+      end
+
+      def create(params = {})
+        slice = new(params)
+        slice.save
+        slice
+      end
+
+      def create!(params = {})
+        slice = new(params)
+        slice.save!
+        slice
+      end
+
+      # Returns output slices in the order of their id
+      # which is usually the order in which they were written.
+      def each
+        all.sort(id: 1).each { |document| yield(document) }
+      end
+
+      # Insert a new slice into the collection
+      #
+      # Returns [Integer] the number of records uploaded
+      #
+      # Parameters
+      #   slice [RocketJob::Sliced::Slice | Array]
+      #     The slice to write to the slices collection
+      #     If slice is an Array, it will be converted to a Slice before inserting
+      #     into the slices collection
+      #
+      #   input_slice [RocketJob::Sliced::Slice]
+      #     The input slice to which this slice corresponds
+      #     The id of the input slice is copied across
+      #     If the insert results in a duplicate record it is ignored, to support
+      #     restarting of jobs that failed in the middle of processing.
+      #     A warning is logged that the slice has already been processed.
+      #
+      # Note:
+      #   `slice_size` is not enforced.
+      #   However many records are present in the slice will be written as a
+      #   single slice to the slices collection
+      #
+      def insert(slice, input_slice = nil)
+        slice = new(records: slice) unless slice.is_a?(Slice)
+
+        # Retain input_slice id in the new output slice
+        if input_slice
+          slice.id                  = input_slice.id
+          slice.first_record_number = input_slice.first_record_number
+        end
+
+        begin
+          slice.save!
+        rescue Mongo::Error::OperationFailure => exc
+          # Ignore duplicates since it means the job was restarted
+          raise(exc) unless exc.message.include?('E11000')
+          logger.warn "Skipped already processed slice# #{slice.id}"
+        end
+        slice
+      end
+
+      alias << insert
+
+      # Index for find_and_modify only if it is not already present
+      def create_indexes
+        all.collection.indexes.create_one(state: 1, _id: 1) if all.collection.indexes.none? { |i| i['name'] == 'state_1__id_1' }
+      rescue Mongo::Error::OperationFailure
+        all.collection.indexes.create_one(state: 1, _id: 1)
+      end
+
+      def_instance_delegators :@all, :collection, :count, :delete_all, :first, :find, :last, :nor, :not, :or, :to_a, :where
+
+      # Drop this collection when it is no longer needed
+      def drop
+        all.collection.drop
+      end
+
+      # Forwardable generates invalid warnings on these methods.
+      def completed
+        all.completed
+      end
+
+      def failed
+        all.failed
+      end
+
+      def queued
+        all.queued
+      end
+
+      def running
+        all.running
+      end
+
+      # Mongoid does not apply ordering, add sort
+      def first
+        all.sort('_id' => 1).first
+      end
+
+      def last
+        all.sort('_id' => -1).first
+      end
+
+      # Returns [Array<Struct>] grouped exceptions by class name,
+      # and unique exception messages by exception class.
+      #
+      # Each struct consists of:
+      #   class_name: [String]
+      #     Exception class name.
+      #
+      #   count: [Integer]
+      #     Number of exceptions with this class.
+      #
+      #   messages: [Array<String>]
+      #     Unique list of error messages.
+      def group_exceptions
+        result_struct = Struct.new(:class_name, :count, :messages)
+        result        = all.collection.aggregate(
+          [
+            {
+              '$match' => {state: 'failed'}
+            },
+            {
+              '$group' => {
+                _id:      {error_class: '$exception.class_name'},
+                messages: {'$addToSet' => '$exception.message'},
+                count:    {'$sum' => 1}
+              }
+            }
+          ]
+        )
+        result.collect do |errors|
+          result_struct.new(errors['_id']['error_class'], errors['count'], errors['messages'])
+        end
+      end
+    end
+  end
+end

data/lib/rocket_job/sliced/writer/input.rb

@@ -0,0 +1,60 @@
+module RocketJob
+  module Sliced
+    module Writer
+      # Internal class for uploading records into input slices
+      class Input
+        attr_reader :record_count
+
+        # Batch collection of lines into slices.
+        #
+        # Parameters
+        #   on_first: [Proc]
+        #     Block to call on the first line only, instead of storing in the slice.
+        #     Useful for extracting the header row
+        #     Default: nil
+        def self.collect(input, **args, &block)
+          writer = new(input, **args)
+          block.call(writer)
+          writer.record_count
+        rescue Exception => exc
+          # Drop input collection when upload fails
+          input.drop
+          raise exc
+        ensure
+          writer&.close
+        end
+
+        def initialize(input, on_first: nil)
+          @on_first      = on_first
+          @batch_count   = 0
+          @record_count  = 0
+          @input         = input
+          @record_number = 1
+          @slice         = @input.new(first_record_number: @record_number)
+        end
+
+        def <<(line)
+          @record_number += 1
+          if @on_first
+            @on_first.call(line)
+            @on_first = nil
+            return self
+          end
+          @slice << line
+          @batch_count  += 1
+          @record_count += 1
+          if @batch_count >= @input.slice_size
+            @input.insert(@slice)
+            @batch_count = 0
+            @slice       = @input.new(first_record_number: @record_number)
+          end
+          self
+        end
+
+        def close
+          @input.insert(@slice) if @slice.size.positive?
+        end
+      end
+    end
+  end
+end