Add support for find_in_batches although not sure it really makes muc…

…h sense for CPK. Fixes #126

lib/composite_primary_keys/relation.rb

@@ -2,6 +2,7 @@ module ActiveRecord
   class Relation
     def add_cpk_support
       class << self
+        include CompositePrimaryKeys::ActiveRecord::Batches
         include CompositePrimaryKeys::ActiveRecord::Calculations
         include CompositePrimaryKeys::ActiveRecord::FinderMethods
         include CompositePrimaryKeys::ActiveRecord::QueryMethods

lib/composite_primary_keys/relation/batches.rb

@@ -0,0 +1,63 @@
+module CompositePrimaryKeys
+  module ActiveRecord
+    module Batches
+      def find_in_batches(options = {})
+        relation = self
+
+        unless arel.orders.blank? && arel.taken.blank?
+          ActiveRecord::Base.logger.warn("Scoped order and limit are ignored, it's forced to be batch order and batch size")
+        end
+
+        if (finder_options = options.except(:start, :batch_size)).present?
+          raise "You can't specify an order, it's forced to be #{batch_order}" if options[:order].present?
+          raise "You can't specify a limit, it's forced to be the batch_size"  if options[:limit].present?
+
+          relation = apply_finder_options(finder_options)
+        end
+
+        start = options.delete(:start).to_i
+        batch_size = options.delete(:batch_size) || 1000
+
+        relation = relation.reorder(batch_order).limit(batch_size)
+
+        # CPK
+        #records = relation.where(table[primary_key].gteq(start)).all
+        self.primary_key.each do |key|
+          relation = relation.where(table[key].gteq(start))
+        end
+        records = relation.all
+
+        while records.any?
+          records_size = records.size
+          primary_key_offset = records.last.id
+
+          yield records
+
+          break if records_size < batch_size
+
+          if primary_key_offset
+            # CPK
+            #records = relation.where(table[primary_key].gt(primary_key_offset)).to_a
+            self.primary_key.each do |key|
+              relation = relation.where(table[key].gt(primary_key_offset))
+            end
+            records = relation.to_a
+
+          else
+            raise "Primary key not included in the custom select clause"
+          end
+        end
+      end
+
+      private
+
+      def batch_order
+        # CPK
+        #"#{quoted_table_name}.#{quoted_primary_key} ASC"
+        self.primary_key.map do |key|
+          "#{quoted_table_name}.#{key} ASC"
+        end.join(",")
+      end
+    end
+  end
+end

test/test_find.rb

@@ -2,7 +2,7 @@
 
 # Testing the find action on composite ActiveRecords with two primary keys
 class TestFind < ActiveSupport::TestCase
-  fixtures :capitols, :reference_types, :reference_codes, :suburbs
+  fixtures :capitols, :departments, :reference_types, :reference_codes, :suburbs
 
   def test_find_first
     ref_code = ReferenceCode.find(:first, :order => 'reference_type_id, reference_code')
@@ -75,4 +75,10 @@ def test_find_last_suburb_with_order
     suburb = Suburb.find(:last, :order => 'suburbs.city_id DESC')
     assert_equal([1,1], suburb.id)
   end
+
+  def test_find_in_batchs
+    departments = Department.find_in_batches do |batch|
+      assert_equal(2, batch.size)
+    end
+  end
 end