Improve code from 231d7676f72947bae765b9bd885b134aaf949921

1 files changed, 5 insertions, 7 deletions

diff --git a/activerecord/lib/active_record/relation/batches.rb b/activerecord/lib/active_record/relation/batches.rb
index 4649e3b376..412be895c4 100644
--- a/activerecord/lib/active_record/relation/batches.rb
+++ b/activerecord/lib/active_record/relation/batches.rb
@@ -67,18 +67,16 @@ module ActiveRecord
       relation = relation.except(:order).order(batch_order).limit(batch_size)
       records = relation.where(primary_key.gteq(start)).all
 
-      key_value = self.primary_key.name
-
       while records.any?
         yield records
 
         break if records.size < batch_size
 
-        last_value = records.last.send(key_value)
-
-        raise "You must include the primary key if you define a select" unless last_value.present?
-
-        records = relation.where(primary_key.gt(last_value)).all
+        if primary_key_offset = records.last.id
+          records = relation.where(primary_key.gt(primary_key_offset)).all
+        else
+          raise "Primary key not included in the custom select clause"
+        end
       end
     end