From 74f8d30a0e4049400344e29fde70587d96e7b834 Mon Sep 17 00:00:00 2001
From: Adrianna Chang <adrianna.chang@shopify.com>
Date: Mon, 31 May 2021 14:14:58 -0400
Subject: [PATCH] Bring in patch from Rails to expose public attr readers on
 BatchEnumerator

---
 .../maintenance_tasks/task_job_concern.rb     |  9 +++-----
 lib/maintenance_tasks.rb                      |  2 ++
 lib/patches/active_record_batch_enumerator.rb | 23 +++++++++++++++++++
 3 files changed, 28 insertions(+), 6 deletions(-)
 create mode 100644 lib/patches/active_record_batch_enumerator.rb

diff --git a/app/jobs/concerns/maintenance_tasks/task_job_concern.rb b/app/jobs/concerns/maintenance_tasks/task_job_concern.rb
index 73549097c..8a8fc312d 100644
--- a/app/jobs/concerns/maintenance_tasks/task_job_concern.rb
+++ b/app/jobs/concerns/maintenance_tasks/task_job_concern.rb
@@ -37,17 +37,14 @@ def build_enumerator(_run, cursor:)
       when ActiveRecord::Relation
         enumerator_builder.active_record_on_records(collection, cursor: cursor)
       when ActiveRecord::Batches::BatchEnumerator
-        if collection.instance_variable_get(:@start) ||
-          collection.instance_variable_get(:@finish)
+        if collection.start || collection.finish
           raise ArgumentError, "#{@task.class.name}#collection cannot support "\
             "a batch enumerator with the +start+ or +finish+ options."
         end
-        relation = collection.instance_variable_get(:@relation)
-        batch_size = collection.instance_variable_get(:@of)
         enumerator_builder.active_record_on_batch_relations(
-          relation,
+          collection.relation,
           cursor: cursor,
-          batch_size: batch_size,
+          batch_size: collection.batch_size,
         )
       when Array
         enumerator_builder.build_array_enumerator(collection, cursor: cursor)
diff --git a/lib/maintenance_tasks.rb b/lib/maintenance_tasks.rb
index fda5157e4..ff32c57d9 100644
--- a/lib/maintenance_tasks.rb
+++ b/lib/maintenance_tasks.rb
@@ -7,6 +7,8 @@
 require "job-iteration"
 require "maintenance_tasks/engine"
 
+require "patches/active_record_batch_enumerator"
+
 # The engine's namespace module. It provides isolation between the host
 # application's code and the engine-specific code. Top-level engine constants
 # and variables are defined under this module.
diff --git a/lib/patches/active_record_batch_enumerator.rb b/lib/patches/active_record_batch_enumerator.rb
new file mode 100644
index 000000000..72a891e12
--- /dev/null
+++ b/lib/patches/active_record_batch_enumerator.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+# TODO: Remove this patch once https://github.com/rails/rails/pull/42312/commits/a031a43d969c87542c4ee8d0d338d55fcbb53376
+# is released.
+module ActiveRecordBatchEnumerator
+  # The primary key value from which the BatchEnumerator starts,
+  #   inclusive of the value.
+  attr_reader :start
+
+  # The primary key value at which the BatchEnumerator ends,
+  #   inclusive of the value.
+  attr_reader :finish
+
+  # The relation from which the BatchEnumerator yields batches.
+  attr_reader :relation
+
+  # The size of the batches yielded by the BatchEnumerator.
+  def batch_size
+    @of
+  end
+end
+
+ActiveRecord::Batches::BatchEnumerator.include(ActiveRecordBatchEnumerator)