Merge pull request #431 from Shopify/ac-fix-cpk-enumerators

Support AR Enumerators for models with composite primary keys using :id

Author: adrianna-chang-shopify

lib/job-iteration/active_record_cursor.rb

@@ -55,7 +55,12 @@ def position=(position)
     def update_from_record(record)
       self.position = @columns.map do |column|
         method = column.to_s.split(".").last
-        record.send(method.to_sym)
+
+        if ActiveRecord.version >= Gem::Version.new("7.1.0.alpha") && method == "id"
+          record.id_value
+        else
+          record.send(method.to_sym)
+        end
       end
     end
 

test/test_helper.rb

@@ -53,8 +53,8 @@ class TravelRoute < ActiveRecord::Base
   self.primary_key = [:origin, :destination]
 end
 
-class TravelRoute < ActiveRecord::Base
-  self.primary_key = [:origin, :destination]
+class Order < ActiveRecord::Base
+  self.primary_key = [:shop_id, :id]
 end
 
 host = ENV["USING_DEV"] == "1" ? "job-iteration.railgun" : "localhost"
@@ -96,6 +96,11 @@ class TravelRoute < ActiveRecord::Base
     t.string(:destination)
     t.string(:origin)
   end
+
+  create_table(:orders, force: true) do |t|
+    t.integer(:shop_id)
+    t.string(:name)
+  end
 end
 
 module LoggingHelpers

test/unit/active_record_enumerator_test.rb

@@ -105,18 +105,33 @@ class ActiveRecordEnumeratorTest < IterationUnitTest
       assert_equal(10, enum.size)
     end
 
-    test "enumerator for a relation with a composite primary key" do
-      TravelRoute.create!(origin: "A", destination: "B")
-      TravelRoute.create!(origin: "A", destination: "C")
-      TravelRoute.create!(origin: "B", destination: "A")
+    if ActiveRecord.version >= Gem::Version.new("7.1.0.alpha")
+      test "enumerator for a relation with a composite primary key" do
+        TravelRoute.create!(origin: "A", destination: "B")
+        TravelRoute.create!(origin: "A", destination: "C")
+        TravelRoute.create!(origin: "B", destination: "A")
 
-      enum = build_enumerator(relation: TravelRoute.all, batch_size: 2)
+        enum = build_enumerator(relation: TravelRoute.all, batch_size: 2)
 
-      cursors = []
-      enum.records.each { |_record, cursor| cursors << cursor }
+        cursors = []
+        enum.records.each { |_record, cursor| cursors << cursor }
 
-      assert_equal([["A", "B"], ["A", "C"], ["B", "A"]], cursors)
-    end if ActiveRecord.version >= Gem::Version.new("7.1.0.alpha")
+        assert_equal([["A", "B"], ["A", "C"], ["B", "A"]], cursors)
+      end
+
+      test "enumerator for a relation with a composite primary key using :id" do
+        Order.create!(name: "Yellow socks", shop_id: 3)
+        Order.create!(name: "Red hat", shop_id: 1)
+        Order.create!(name: "Blue jeans", shop_id: 1)
+
+        enum = build_enumerator(relation: Order.all, batch_size: 2)
+
+        order_names = []
+        enum.records.each { |record, _cursor| order_names << record.name }
+
+        assert_equal(["Red hat", "Blue jeans", "Yellow socks"], order_names)
+      end
+    end
 
     private