sidekiq-iteration 0.3.0 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 40efca13e06cd7fdcfc1ff59ad08fea8fc731ee1b5560ae5b75b2591379bcb63
-  data.tar.gz: eec2991b40bb67ffc1dcea55f1c0e8acc98a18cd59ad2fd117cd80c1a94c3e79
+  metadata.gz: ea1fc3e6f5faff037ecfade45cc58bd2035385cafb888ef67ce253d12295aee8
+  data.tar.gz: a4bf097d4a1a8750f4d3e2ac9ce4f01191b5e7313635fed3d958b89647639c9a
 SHA512:
-  metadata.gz: 1162ffafc4d157e7a8f9d2b8f69163e90e83431daac129707e16605a9b0df250c1cf2dd9063f651782b01ab0dcd9a0f3848d381981f94ef5a2daf36f43d591be
-  data.tar.gz: 45a1efa4e1e65ae322b7c923cb5795ceda901863afc4d135cb25e1b342c9604e7f90719259b7fd93b8c38b15a5bcb4cab984617ea915fb58055f21936470269c
+  metadata.gz: 415f27277011c3721853ae64c1c78c566a3bfe2c690ebaeb8e9e05bb10b4f3faac00cda0f6f95ef9fd980993695edc6621085833c7901bb09b14ff6027467622
+  data.tar.gz: 68568c8205c0370a3d1765e5bd6431655d1a3d5fd0ff07ee65fc6e8bf1dd0447fba1a62978c45fea3850029a9046f4cbe7768bce885f5bcb4d9c8e322a539ab6

data/CHANGELOG.md

@@ -1,5 +1,35 @@
 ## master (unreleased)
 
+## 0.4.0 (2024-05-10)
+
+- Support ordering using multiple directions for ActiveRecord enumerators
+
+    ```ruby
+    active_record_records_enumerator(..., columns: [:shop_id, :id], order: [:asc, :desc])
+    ```
+
+- Support iterating over ActiveRecord models with composite primary keys
+
+- Use Arel to generate SQL in ActiveRecord enumerator
+
+    Previously, the enumerator coerced numeric ids to a string value (e.g.: `... AND id > '1'`),
+    which can cause problems on some DBMSes (like BigQuery).
+
+- Enforce explicitly passed to ActiveRecord enumerators `:columns` value to include a primary key
+
+    Previously, the primary key column was added implicitly if it was not in the list.
+
+    ```ruby
+    # before
+    active_record_records_enumerator(..., columns: [:updated_at])
+
+    # after
+    active_record_records_enumerator(..., columns: [:updated_at, :id])
+    ```
+
+- Accept single values as a `:columns` for ActiveRecord enumerators
+- Add `around_iteration` hook
+
 ## 0.3.0 (2023-05-20)
 
 - Allow a default retry backoff to be configured

data/README.md

@@ -4,6 +4,8 @@
 
 Meet Iteration, an extension for [Sidekiq](https://github.com/mperham/sidekiq) that makes your long-running jobs interruptible and resumable, saving all progress that the job has made (aka checkpoint for jobs).
 
+You may consider [`pluck_in_batches`](https://github.com/fatkodima/pluck_in_batches) gem to speedup iterating over large database tables.
+
 ## Background
 
 Imagine the following job:
@@ -99,6 +101,12 @@ class NotifyUsersJob
     # Will be called when the job starts iterating. Called only once, for the first time.
   end
 
+  def around_iteration
+    # Will be called around each iteration.
+    # Can be useful for some metrics collection, performance tracking etc.
+    yield
+  end
+
   def on_resume
     # Called when the job resumes iterating.
   end

data/guides/best-practices.md

@@ -2,7 +2,7 @@
 
 ## Considerations when writing jobs
 
-* Duration of `#each_iteration`: processing a single element from the enumerator builded in `#build_enumerator` should take less than 25 seconds, or the duration set as a timeout for Sidekiq. It allows the job to be safely interrupted and resumed.
+* Duration of `#each_iteration`: processing a single element from the enumerator built in `#build_enumerator` should take less than 25 seconds, or the duration set as a timeout for Sidekiq. It allows the job to be safely interrupted and resumed.
 * Idempotency of `#each_iteration`: it should be safe to run `#each_iteration` multiple times for the same element from the enumerator. Read more in [this Sidekiq best practice](https://github.com/mperham/sidekiq/wiki/Best-Practices#2-make-your-job-idempotent-and-transactional). It's important if the job errors and you run it again, because the same element that errored the job may be processed again. It especially matters in the situation described above, when the iteration duration exceeds the timeout: if the job is re-enqueued, multiple elements may be processed again.
 
 ## Batch iteration

data/guides/custom-enumerator.md

@@ -6,7 +6,7 @@ Before writing an enumerator, it is important to understand [how Iteration works
 your enumerator will be used by it. An enumerator must `yield` two things in the following order as positional
 arguments:
 - An object to be processed in a job `each_iteration` method
-- A cursor position, which Iteration will persist if `each_iteration` returns succesfully and the job is forced to shut
+- A cursor position, which Iteration will persist if `each_iteration` returns successfully and the job is forced to shut
   down. It can be any data type your job backend can serialize and deserialize correctly.
 
 A job that includes Iteration is first started with `nil` as the cursor. When resuming an interrupted job, Iteration

data/guides/iteration-how-it-works.md

@@ -36,6 +36,12 @@ SELECT "users".* FROM "users" ORDER BY "users"."id" LIMIT 100
 SELECT "users".* FROM "users" WHERE "users"."id" > 2 ORDER BY "products"."id" LIMIT 100
 ```
 
+## Exceptions inside `each_iteration`
+
+When an unrescued exception happens inside the `each_iteration` block, the job will stop and re-enqueue itself with the last successful cursor. This means that the iteration that failed will be retried with the same parameters and the cursor will only move if that iteration succeeds. This behaviour may be enough for intermittent errors, such as network connection failures, but if your execution is deterministic and you have an error, subsequent iterations will never run.
+
+In other words, if you are trying to process 100 records but the job consistently fails on the 61st, only the first 60 will be processed and the job will try to process the 61st record until retries are exhausted.
+
 ## Signals
 
 It's critical to know [UNIX signals](https://www.tutorialspoint.com/unix/unix-signals-traps.htm) in order to understand how interruption works. There are two main signals that Sidekiq use: `SIGTERM` and `SIGKILL`. `SIGTERM` is the graceful termination signal which means that the process should exit _soon_, not immediately. For Iteration, it means that we have time to wait for the last iteration to finish and to push job back to the queue with the last cursor position.

data/guides/throttling.md

@@ -25,7 +25,7 @@ class DeleteAccountsThrottledJob
 end
 ```
 
-Note that it’s up to you to define a throttling condition that makes sense for your app.
+Note that it's up to you to define a throttling condition that makes sense for your app.
 For example, `DatabaseStatus.healthy?` can check various MySQL metrics such as replication lag, DB threads, whether DB writes are available, etc.
 
 Jobs can define multiple throttle conditions. Throttle conditions are inherited by descendants, and new conditions will be appended without impacting existing conditions.

data/lib/sidekiq_iteration/active_record_enumerator.rb

@@ -10,42 +10,64 @@ module SidekiqIteration
         raise ArgumentError, "relation must be an ActiveRecord::Relation"
       end
 
-      unless order == :asc || order == :desc
-        raise ArgumentError, ":order must be :asc or :desc, got #{order.inspect}"
-      end
-
-      @primary_key = "#{relation.table_name}.#{relation.primary_key}"
-      @columns = Array(columns&.map(&:to_s) || @primary_key)
-      @primary_key_index = @columns.index(@primary_key) || @columns.index(relation.primary_key)
-      @pluck_columns = if @primary_key_index
-                         @columns
-                       else
-                         @columns + [@primary_key]
-                       end
-      @batch_size = batch_size
-      @order = order
-      @cursor = Array.wrap(cursor)
-      raise ArgumentError, "Must specify at least one column" if @columns.empty?
-      if relation.joins_values.present? && !@columns.all?(/\./)
-        raise ArgumentError, "You need to specify fully-qualified columns if you join a table"
-      end
-
       if relation.arel.orders.present? || relation.arel.taken.present?
         raise ArgumentError,
           "The relation cannot use ORDER BY or LIMIT due to the way how iteration with a cursor is designed. " \
           "You can use other ways to limit the number of rows, e.g. a WHERE condition on the primary key column."
       end
 
-      ordering = @columns.to_h { |column| [column, @order] }
+      @relation = relation
+      @primary_key = relation.primary_key
+      columns = Array(columns || @primary_key).map(&:to_s)
+
+      if (Array(order) - [:asc, :desc]).any?
+        raise ArgumentError, ":order must be :asc or :desc or an array consisting of :asc or :desc, got #{order.inspect}"
+      end
+
+      if order.is_a?(Array) && order.size != columns.size
+        raise ArgumentError, ":order must include a direction for each batching column"
+      end
+
+      @primary_key_index = primary_key_index(columns, relation)
+      if @primary_key_index.nil? || (composite_primary_key? && @primary_key_index.any?(nil))
+        raise ArgumentError, ":columns must include a primary key columns"
+      end
+
+      @batch_size = batch_size
+      @order = batch_order(columns, order)
+      @cursor = Array(cursor)
+
+      if @cursor.present? && @cursor.size != columns.size
+        raise ArgumentError, ":cursor must include values for all the columns from :columns"
+      end
+
+      if columns.any?(/\W/)
+        arel_columns = columns.map.with_index do |column, i|
+          arel_column(column).as("cursor_column_#{i + 1}")
+        end
+        @cursor_columns = arel_columns.map { |column| column.right.to_s }
+
+        relation =
+          if relation.select_values.empty?
+            relation.select(@relation.arel_table[Arel.star], arel_columns)
+          else
+            relation.select(arel_columns)
+          end
+      else
+        @cursor_columns = columns
+      end
+
+      @columns = columns
+      ordering = @columns.zip(@order).to_h
       @base_relation = relation.reorder(ordering)
       @iteration_count = 0
     end
 
     def records
       Enumerator.new(-> { records_size }) do |yielder|
-        batches.each do |batch, _|
+        batches.each do |batch, _| # rubocop:disable Style/HashEachMethods
           batch.each do |record|
-            @iteration_count += 1
+            increment_iteration
             yielder.yield(record, cursor_value(record))
           end
         end
@@ -55,7 +77,7 @@ module SidekiqIteration
     def batches
       Enumerator.new(-> { records_size }) do |yielder|
         while (batch = next_batch(load: true))
-          @iteration_count += 1
+          increment_iteration
           yielder.yield(batch, cursor_value(batch.last))
         end
       end
@@ -64,13 +86,44 @@ module SidekiqIteration
     def relations
       Enumerator.new(-> { relations_size }) do |yielder|
         while (batch = next_batch(load: false))
-          @iteration_count += 1
+          increment_iteration
           yielder.yield(batch, unwrap_array(@cursor))
         end
       end
     end
 
     private
+      def primary_key_index(columns, relation)
+        indexes = Array(@primary_key).map do |pk_column|
+          columns.index do |column|
+            column == pk_column ||
+              (column.include?(relation.table_name) && column.include?(pk_column))
+          end
+        end
+
+        if composite_primary_key?
+          indexes
+        else
+          indexes.first
+        end
+      end
+
+      def batch_order(columns, order)
+        if order.is_a?(Array)
+          order
+        else
+          [order] * columns.size
+        end
+      end
+
+      def arel_column(column)
+        if column.include?(".")
+          Arel.sql(column)
+        else
+          @relation.arel_table[column]
+        end
+      end
+
       def records_size
         @base_relation.count(:all)
       end
@@ -81,8 +134,8 @@ module SidekiqIteration
 
       def next_batch(load:)
         batch_relation = @base_relation.limit(@batch_size)
-        if conditions.any?
-          batch_relation = batch_relation.where(*conditions)
+        if @cursor.present?
+          batch_relation = apply_cursor(batch_relation)
         end
 
         records = nil
@@ -98,9 +151,7 @@ module SidekiqIteration
         cursor = cursor_values.last
         return unless cursor.present?
 
-        # The primary key was plucked, but original cursor did not include it, so we should remove it
-        cursor.pop unless @primary_key_index
-        @cursor = Array.wrap(cursor)
+        @cursor = Array(cursor)
 
         # Yields relations by selecting the primary keys of records in the batch.
         # Post.where(published: nil) results in an enumerator of relations like:
@@ -111,80 +162,89 @@ module SidekiqIteration
       end
 
       def pluck_columns(batch)
-        columns =
-          if batch.is_a?(Array)
-            @pluck_columns.map { |column| column.to_s.split(".").last }
-          else
-            @pluck_columns
-          end
-
-        if columns.size == 1 # only the primary key
-          column_values = batch.pluck(columns.first)
+        if @cursor_columns.size == 1 # only the primary key
+          column_values = batch.pluck(@cursor_columns.first)
           return [column_values, column_values]
         end
 
-        column_values = batch.pluck(*columns)
-        primary_key_index = @primary_key_index || -1
-        primary_key_values = column_values.map { |values| values[primary_key_index] }
+        column_values = batch.pluck(*@cursor_columns)
+        primary_key_values =
+          if composite_primary_key?
+            column_values.map { |values| values.values_at(*@primary_key_index) }
+          else
+            column_values.map { |values| values[@primary_key_index] }
+          end
 
-        serialize_column_values!(column_values)
+        column_values = serialize_column_values(column_values)
         [column_values, primary_key_values]
       end
 
       def cursor_value(record)
-        positions = @columns.map do |column|
-          attribute_name = column.to_s.split(".").last
-          column_value(record[attribute_name])
+        positions = @cursor_columns.map do |column|
+          column_value(record[column])
         end
 
         unwrap_array(positions)
       end
 
-      def conditions
-        return [] if @cursor.empty?
+      # (x, y) >= (a, b) iff (x > a or (x = a and y >= b))
+      # (x, y) <= (a, b) iff (x < a or (x = a and y <= b))
+      def apply_cursor(relation)
+        arel_columns = @columns.map { |column| arel_column(column) }
+        cursor_positions = arel_columns.zip(@cursor, cursor_operators)
 
-        binds = []
-        sql = build_starts_after_conditions(0, binds)
-
-        # Start from the record pointed by cursor.
-        # We use the property that `>=` is equivalent to `> or =`.
-        if @iteration_count == 0
-          binds.unshift(*@cursor)
-          columns_equality = @columns.map { |column| "#{column} = ?" }.join(" AND ")
-          sql = "(#{columns_equality}) OR (#{sql})"
+        where_clause = nil
+        cursor_positions.reverse_each.with_index do |(arel_column, value, operator), index|
+          where_clause =
+            if index == 0
+              arel_column.public_send(operator, value)
+            else
+              arel_column.public_send(operator, value).or(
+                arel_column.eq(value).and(where_clause),
+              )
+            end
         end
 
-        [sql, *binds]
+        relation.where(where_clause)
       end
 
-      # (x, y) > (a, b) iff (x > a or (x = a and y > b))
-      # (x, y) < (a, b) iff (x < a or (x = a and y < b))
-      def build_starts_after_conditions(index, binds)
-        column = @columns[index]
+      def serialize_column_values(column_values)
+        column_values.map { |values| values.map { |value| column_value(value) } }
+      end
 
-        if index < @cursor.size - 1
-          binds << @cursor[index] << @cursor[index]
-          "#{column} #{@order == :asc ? '>' : '<'} ? OR (#{column} = ? AND (#{build_starts_after_conditions(index + 1, binds)}))"
+      def column_value(value)
+        if value.is_a?(Time)
+          value.strftime(SQL_DATETIME_WITH_NSEC)
         else
-          binds << @cursor[index]
-          if @columns.size == @cursor.size
-            @order == :asc ? "#{column} > ?" : "#{column} < ?"
+          value
+        end
+      end
+
+      def cursor_operators
+        # Start from the record pointed by cursor when just starting.
+        @columns.zip(@order).map do |column, order|
+          if column == @columns.last
+            if order == :asc
+              first_iteration? ? :gteq : :gt
+            else
+              first_iteration? ? :lteq : :lt
+            end
           else
-            @order == :asc ? "#{column} >= ?" : "#{column} <= ?"
+            order == :asc ? :gt : :lt
           end
         end
       end
 
-      def serialize_column_values!(column_values)
-        column_values.map! { |values| values.map! { |value| column_value(value) } }
+      def increment_iteration
+        @iteration_count += 1
       end
 
-      def column_value(value)
-        if value.is_a?(Time)
-          value.strftime(SQL_DATETIME_WITH_NSEC)
-        else
-          value
-        end
+      def first_iteration?
+        @iteration_count == 0
+      end
+
+      def composite_primary_key?
+        @primary_key.is_a?(Array)
       end
 
       def unwrap_array(array)

data/lib/sidekiq_iteration/csv_enumerator.rb

@@ -36,7 +36,7 @@ module SidekiqIteration
     #   SidekiqIteration::CsvEnumerator.new(csv).rows(cursor: cursor)
     #
     def initialize(csv)
-      unless csv.instance_of?(CSV)
+      unless defined?(CSV) && csv.instance_of?(CSV)
         raise ArgumentError, "CsvEnumerator.new takes CSV object"
       end
 

data/lib/sidekiq_iteration/enumerators.rb

@@ -17,10 +17,6 @@ module SidekiqIteration
     def array_enumerator(array, cursor:)
       raise ArgumentError, "array must be an Array" unless array.is_a?(Array)
 
-      if defined?(ActiveRecord) && array.any?(ActiveRecord::Base)
-        raise ArgumentError, "array cannot contain ActiveRecord objects"
-      end
-
       array.each_with_index.drop(cursor || 0).to_enum { array.size }
     end
 
@@ -28,10 +24,10 @@ module SidekiqIteration
     #
     # @param scope [ActiveRecord::Relation] scope to iterate
     # @param cursor [Object] offset to start iteration from, usually an id
-    # @option options :columns [Array<String, Symbol>] used to build the actual query for iteration,
+    # @option options :columns [Array<String, Symbol>, String, Symbol] used to build the actual query for iteration,
     #   defaults to primary key
     # @option options :batch_size [Integer] (100) size of the batch
-    # @option options :order [:asc, :desc] (:asc) specifies iteration order
+    # @option options :order [:asc, :desc, Array<:asc, :desc>] (:asc) specifies iteration order
     #
     # +columns:+ argument is used to build the actual query for iteration. +columns+: defaults to primary key:
     #
@@ -59,7 +55,7 @@ module SidekiqIteration
     # As a result of this query pattern, if the values in these columns change for the records in scope during
     # iteration, they may be skipped or yielded multiple times depending on the nature of the update and the
     # cursor's value. If the value gets updated to a greater value than the cursor's value, it will get yielded
-    # again. Similarly, if the value gets updated to a lesser value than the curor's value, it will get skipped.
+    # again. Similarly, if the value gets updated to a lesser value than the cursor's value, it will get skipped.
     #
     # @example
     #   def build_enumerator(cursor:)

data/lib/sidekiq_iteration/iteration.rb

@@ -20,8 +20,7 @@ module SidekiqIteration
         end
 
         throttle_on(backoff: SidekiqIteration.default_retry_backoff) do
-          defined?(Sidekiq::CLI) &&
-            Sidekiq::CLI.instance.launcher.stopping?
+          SidekiqIteration.stopping
         end
       end
 
@@ -88,6 +87,12 @@ module SidekiqIteration
     def on_start
     end
 
+    # A hook to override that will be called around each iteration.
+    # Can be useful for some metrics collection, performance tracking etc.
+    def around_iteration
+      yield
+    end
+
     # A hook to override that will be called when the job resumes iterating.
     def on_resume
     end
@@ -178,7 +183,9 @@ module SidekiqIteration
 
         enumerator.each do |object_from_enumerator, index|
           found_record = true
-          each_iteration(object_from_enumerator, *arguments)
+          around_iteration do
+            each_iteration(object_from_enumerator, *arguments)
+          end
           @cursor_position = index
           @current_run_iterations += 1
 
@@ -258,13 +265,6 @@ module SidekiqIteration
           true
         when false, :skip_complete_callback
           false
-        when Array # can be used to return early from the enumerator
-          reason, backoff = completed
-          raise "Unknown reason: #{reason}" unless reason == :retry
-
-          @job_iteration_retry_backoff = backoff
-          @needs_reenqueue = true
-          false
         else
           raise "Unexpected thrown value: #{completed.inspect}"
         end

data/lib/sidekiq_iteration/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module SidekiqIteration
-  VERSION = "0.3.0"
+  VERSION = "0.4.0"
 end

data/lib/sidekiq_iteration.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 require "sidekiq"
+require_relative "sidekiq_iteration/iteration"
+require_relative "sidekiq_iteration/job_retry_patch"
 require_relative "sidekiq_iteration/version"
 
 module SidekiqIteration
@@ -44,8 +46,14 @@ module SidekiqIteration
     def logger
       @logger ||= Sidekiq.logger
     end
+
+    # @private
+    attr_accessor :stopping
   end
 end
 
-require_relative "sidekiq_iteration/iteration"
-require_relative "sidekiq_iteration/job_retry_patch"
+Sidekiq.configure_server do |config|
+  config.on(:quiet) do
+    SidekiqIteration.stopping = true
+  end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: sidekiq-iteration
 version: !ruby/object:Gem::Version
-  version: 0.3.0
+  version: 0.4.0
 platform: ruby
 authors:
 - fatkodima
@@ -9,7 +9,7 @@ authors:
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2023-05-20 00:00:00.000000000 Z
+date: 2024-05-10 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: sidekiq
@@ -72,7 +72,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.4.12
+rubygems_version: 3.4.19
 signing_key:
 specification_version: 4
 summary: Makes your long-running sidekiq jobs interruptible and resumable.