Permalink
Browse files

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

…h sense for CPK. Fixes #126
  • Loading branch information...
1 parent 005b25d commit 97e6a6049e351c3d334de9eed5b0f80496c00df9 @cfis cfis committed Nov 22, 2012
Showing with 71 additions and 1 deletion.
  1. +1 −0 lib/composite_primary_keys/relation.rb
  2. +63 −0 lib/composite_primary_keys/relation/batches.rb
  3. +7 −1 test/test_find.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
@@ -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
View
@@ -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

0 comments on commit 97e6a60

Please sign in to comment.