corrected AR find_each and find_in_batches to raise when the user uses select but does not specify the primary key

Signed-off-by: José Valim <jose.valim@gmail.com>

2 files changed, 22 insertions, 1 deletions

diff --git a/activerecord/lib/active_record/relation/batches.rb b/activerecord/lib/active_record/relation/batches.rb
index 1c61e7d450..4649e3b376 100644
--- a/activerecord/lib/active_record/relation/batches.rb
+++ b/activerecord/lib/active_record/relation/batches.rb
@@ -67,11 +67,18 @@ 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
-        records = relation.where(primary_key.gt(records.last.id)).all
+
+        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
       end
     end
 
diff --git a/activerecord/test/cases/batches_test.rb b/activerecord/test/cases/batches_test.rb
index 83deabb5b7..dcc49e12ca 100644
--- a/activerecord/test/cases/batches_test.rb
+++ b/activerecord/test/cases/batches_test.rb
@@ -17,6 +17,20 @@ class EachTest < ActiveRecord::TestCase
     end
   end
 
+  def test_each_should_raise_if_select_is_set_without_id
+    assert_raise(RuntimeError) do
+      Post.find_each(:select => :title, :batch_size => 1) { |post| post }
+    end
+  end
+
+  def test_each_should_execute_if_id_is_in_select
+    assert_queries(4) do
+      Post.find_each(:select => "id, title, type", :batch_size => 2) do |post|
+        assert_kind_of Post, post
+      end
+    end
+  end
+
   def test_each_should_raise_if_the_order_is_set
     assert_raise(RuntimeError) do
       Post.find_each(:order => "title") { |post| post }