Localize patch for AR::Relation#primary_key in Rails 8.0+

lib/activerecord-bitemporal/bitemporal.rb

@@ -112,15 +112,36 @@ def bitemporal_option_storage=(value)
     end
 
     module Relation
-      module Finder
-        def find(*ids)
-          return super if block_given?
-          all.spawn.yield_self { |obj|
-            def obj.primary_key
-              "bitemporal_id"
+      module BitemporalIdAsPrimaryKey # :nodoc:
+        private
+
+        # Generate a method that temporarily changes the primary key to
+        # bitemporal_id for localizing the effect of the change to only the
+        # method specified by `name`.
+        #
+        # DO NOT use this method outside of this module.
+        def use_bitemporal_id_as_primary_key(name) # :nodoc:
+          module_eval <<~RUBY, __FILE__, __LINE__ + 1
+            def #{name}(...)
+              all.spawn.yield_self { |relation|
+                def relation.primary_key
+                  bitemporal_id_key
+                end
+                relation.method(:#{name}).super_method.call(...)
+              }
             end
-            obj.method(:find).super_method.call(*ids)
-          }
+          RUBY
+        end
+      end
+      extend BitemporalIdAsPrimaryKey
+
+      module Finder
+        extend BitemporalIdAsPrimaryKey
+
+        use_bitemporal_id_as_primary_key :find
+
+        if ActiveRecord.version >= Gem::Version.new("8.0.0")
+          use_bitemporal_id_as_primary_key :exists?
         end
 
         def find_at_time!(datetime, *ids)
@@ -136,6 +157,10 @@ def find_at_time(datetime, *ids)
       end
       include Finder
 
+      if ActiveRecord.version >= Gem::Version.new("8.0.0")
+        use_bitemporal_id_as_primary_key :ids
+      end
+
       def build_arel(*)
         ActiveRecord::Bitemporal.with_bitemporal_option(**bitemporal_option) {
           super
@@ -168,8 +193,12 @@ def load
         end
       end
 
-      def primary_key
-        bitemporal_id_key
+      # Use original primary_key for Active Record 8.0+ as much as possible
+      # to avoid issues with patching primary_key of AR::Relation globally.
+      if ActiveRecord.version < Gem::Version.new("8.0.0")
+        def primary_key
+          bitemporal_id_key
+        end
       end
     end
 

lib/activerecord-bitemporal/scope.rb

@@ -160,7 +160,13 @@ module CollectionProxy
     # @see https://github.com/rails/rails/blob/v7.1.3.4/activerecord/lib/active_record/relation/delegation.rb#L117
     delegate :bitemporal_value, :bitemporal_value=, :valid_datetime, :valid_date,
              :transaction_datetime, :bitemporal_option, :bitemporal_option_merge!,
-             :build_arel, :primary_key, to: :scope
+             :build_arel, to: :scope
+
+    if ActiveRecord.version < Gem::Version.new("8.0.0")
+      delegate :primary_key, to: :scope
+    else
+      delegate :ids, :exists?, to: :scope
+    end
   end
 
   module Scope

spec/activerecord-bitemporal/relation/exists_spec.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ActiveRecord::Bitemporal::Relation do
+  describe "#exists?" do
+    def create_employees_with_history
+      Employee
+        .create!([{ name: "Jane" }, { name: "Tom" } ])
+        .each { |e| e.update!(updated_at: Time.current) }
+    end
+
+    # Assuming that the number of records will not reach the maximum value in tests.
+    # employees.bitemporal_id is a 4-byte signed integer.
+    MAX_BITEMPORAL_ID = 1 << 31 - 1
+
+    context "on BTDM" do
+      let(:employees) { create_employees_with_history }
+
+      it "finds existing bitemporal IDs" do
+        first_employee, second_employee = *employees
+
+        expect(Employee.exists?(first_employee.bitemporal_id)).to eq true
+        expect(Employee.exists?(second_employee.bitemporal_id)).to eq true
+        expect(Employee.exists?(MAX_BITEMPORAL_ID)).to eq false
+      end
+    end
+
+    context "on relation" do
+      let(:employees) { create_employees_with_history }
+
+      it "finds existing bitemporal IDs" do
+        first_employee, second_employee = *employees
+
+        expect(Employee.all.exists?(first_employee.bitemporal_id)).to eq true
+        expect(Employee.all.exists?(second_employee.bitemporal_id)).to eq true
+        expect(Employee.all.exists?(MAX_BITEMPORAL_ID)).to eq false
+      end
+    end
+
+    context "on loaded relation" do
+      let(:employees) { create_employees_with_history }
+
+      it "finds existing bitemporal IDs" do
+        first_employee, second_employee = *employees
+
+        expect(Employee.all.load.exists?(first_employee.bitemporal_id)).to eq true
+        expect(Employee.all.load.exists?(second_employee.bitemporal_id)).to eq true
+        expect(Employee.all.load.exists?(MAX_BITEMPORAL_ID)).to eq false
+      end
+    end
+
+    context "with eager loading by includes" do
+      let(:company) {
+        Company
+          .create!(name: "Company")
+          .tap { |c| c.employees << create_employees_with_history }
+      }
+
+      it "finds existing bitemporal IDs" do
+        first_employee, second_employee = *company.employees
+
+        expect(Employee.includes(:company).exists?(first_employee.bitemporal_id)).to eq true
+        expect(Employee.includes(:company).exists?(second_employee.bitemporal_id)).to eq true
+        expect(Employee.includes(:company).exists?(MAX_BITEMPORAL_ID)).to eq false
+      end
+    end
+
+    context "on association" do
+      let(:company) {
+        Company
+          .create!(name: "Company")
+          .tap { |c| c.employees << create_employees_with_history }
+      }
+
+      it "finds existing bitemporal IDs" do
+        first_employee, second_employee = *company.employees
+
+        expect(company.employees.reset.exists?(first_employee.bitemporal_id)).to eq true
+        expect(company.employees.reset.exists?(second_employee.bitemporal_id)).to eq true
+        expect(company.employees.reset.exists?(MAX_BITEMPORAL_ID)).to eq false
+      end
+    end
+  end
+end

spec/activerecord-bitemporal/relation/ids_spec.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ActiveRecord::Bitemporal::Relation do
+  describe "#ids" do
+    def create_employees_with_history
+      Employee
+        .create!([{ name: "Jane" }, { name: "Tom" } ])
+        .each { |e| e.update!(updated_at: Time.current) }
+    end
+
+    context "on BTDM" do
+      let(:employees) { create_employees_with_history }
+
+      it "returns all bitemporal IDs" do
+        expected = [employees[0].bitemporal_id, employees[1].bitemporal_id]
+        expect(Employee.ids).to match_array expected
+      end
+    end
+
+    context "on relation" do
+      let(:employees) { create_employees_with_history }
+
+      it "returns all bitemporal IDs" do
+        expected = [employees[0].bitemporal_id, employees[1].bitemporal_id]
+        expect(Employee.all.ids).to match_array expected
+      end
+    end
+
+    context "on loaded relation" do
+      let(:employees) { create_employees_with_history }
+
+      it "returns all bitemporal IDs" do
+        expected = [employees[0].bitemporal_id, employees[1].bitemporal_id]
+        expect(Employee.all.load.ids).to match_array expected
+      end
+    end
+
+    context "with eager loading by includes" do
+      let(:company) {
+        Company
+          .create!(name: "Company")
+          .tap { |c| c.employees << create_employees_with_history }
+      }
+
+      it "returns bitemporal IDs" do
+        expected = [company.employees[0].bitemporal_id, company.employees[1].bitemporal_id]
+        expect(Employee.includes(:company).ids).to match_array expected
+      end
+    end
+
+    context "on association" do
+      let(:company) {
+        Company
+          .create!(name: "Company")
+          .tap { |c| c.employees << create_employees_with_history }
+      }
+
+      it "returns bitemporal IDs" do
+        expected = [company.employees[0].bitemporal_id, company.employees[1].bitemporal_id]
+        expect(company.employees.reset.ids).to match_array expected
+      end
+    end
+  end
+end