Merge pull request #9969 from divineforest/fix-find-in-batches

Fail early with "Primary key not included in the custom select clause" i...

3 files changed, 13 insertions, 6 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 170818cd1c..3ca689e713 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,11 @@
+*   Fail early with "Primary key not included in the custom select clause"
+    in find_in_batches.
+
+    Before this patch find_in_batches raises this error only on second iteration.
+    So you will know about the problem only when you get the batch size threshold.
+
+    *Alexander Balashov*
+
 *   Ensure `second` through `fifth` methods act like the `first` finder.
 
     The famous ordinal Array instance methods defined in ActiveSupport
diff --git a/activerecord/lib/active_record/relation/batches.rb b/activerecord/lib/active_record/relation/batches.rb
index 49b01909c6..f02e2365f7 100644
--- a/activerecord/lib/active_record/relation/batches.rb
+++ b/activerecord/lib/active_record/relation/batches.rb
@@ -102,16 +102,13 @@ module ActiveRecord
       while records.any?
         records_size = records.size
         primary_key_offset = records.last.id
+        raise "Primary key not included in the custom select clause" unless primary_key_offset
 
         yield records
 
         break if records_size < batch_size
 
-        if primary_key_offset
-          records = relation.where(table[primary_key].gt(primary_key_offset)).to_a
-        else
-          raise "Primary key not included in the custom select clause"
-        end
+        records = relation.where(table[primary_key].gt(primary_key_offset)).to_a
       end
     end
 
diff --git a/activerecord/test/cases/batches_test.rb b/activerecord/test/cases/batches_test.rb
index 38c2560d69..ebb36e4940 100644
--- a/activerecord/test/cases/batches_test.rb
+++ b/activerecord/test/cases/batches_test.rb
@@ -46,7 +46,9 @@ class EachTest < ActiveRecord::TestCase
 
   def test_each_should_raise_if_select_is_set_without_id
     assert_raise(RuntimeError) do
-      Post.select(:title).find_each(:batch_size => 1) { |post| post }
+      Post.select(:title).find_each(batch_size: 1) { |post|
+        flunk "should not call this block"
+      }
     end
   end