ActiveRecordSource: Improve composite primary key support

lib/graphql/dataloader/active_record_source.rb

@@ -7,18 +7,39 @@ class ActiveRecordSource < GraphQL::Dataloader::Source
       def initialize(model_class, find_by: model_class.primary_key)
         @model_class = model_class
         @find_by = find_by
-        @type_for_column = @model_class.type_for_attribute(@find_by)
+        @find_by_many = find_by.is_a?(Array)
+        if @find_by_many
+          @type_for_column = @find_by.map { |fb| @model_class.type_for_attribute(fb) }
+        else
+          @type_for_column = @model_class.type_for_attribute(@find_by)
+        end
       end
 
-      def load(requested_key)
-        casted_key = @type_for_column.cast(requested_key)
-        super(casted_key)
+      def result_key_for(requested_key)
+        normalize_fetch_key(requested_key)
+      end
+
+      def normalize_fetch_key(requested_key)
+        if @find_by_many
+          requested_key.each_with_index.map do |k, idx|
+            @type_for_column[idx].cast(k)
+          end
+        else
+          @type_for_column.cast(requested_key)
+        end
       end
 
       def fetch(record_ids)
         records = @model_class.where(@find_by => record_ids)
         record_lookup = {}
-        records.each { |r| record_lookup[r.public_send(@find_by)] = r }
+        if @find_by_many
+          records.each do |r|
+            key = @find_by.map { |fb| r.public_send(fb) }
+            record_lookup[key] = r
+          end
+        else
+          records.each { |r| record_lookup[r.public_send(@find_by)] = r }
+        end
         record_ids.map { |id| record_lookup[id] }
       end
     end

lib/graphql/dataloader/source.rb

@@ -21,7 +21,7 @@ def setup(dataloader)
       def request(value)
         res_key = result_key_for(value)
         if !@results.key?(res_key)
-          @pending[res_key] ||= value
+          @pending[res_key] ||= normalize_fetch_key(value)
         end
         Dataloader::Request.new(self, value)
       end
@@ -35,12 +35,24 @@ def result_key_for(value)
         value
       end
 
+      # Implement this method if varying values given to {load} (etc) should be consolidated
+      # or normalized before being handed off to your {fetch} implementation.
+      #
+      # This is different than {result_key_for} because _that_ method handles unification inside Dataloader's cache,
+      # but this method changes the value passed into {fetch}.
+      #
+      # @param value [Object] The value passed to {load}, {load_all}, {request}, or {request_all}
+      # @return [Object] The value given to {fetch}
+      def normalize_fetch_key(value)
+        value
+      end
+
       # @return [Dataloader::Request] a pending request for a values from `keys`. Call `.load` on that object to wait for the results.
       def request_all(values)
         values.each do |v|
           res_key = result_key_for(v)
           if !@results.key?(res_key)
-            @pending[res_key] ||= v
+            @pending[res_key] ||= normalize_fetch_key(v)
           end
         end
         Dataloader::RequestAll.new(self, values)
@@ -53,7 +65,7 @@ def load(value)
         if @results.key?(result_key)
           result_for(result_key)
         else
-          @pending[result_key] ||= value
+          @pending[result_key] ||= normalize_fetch_key(value)
           sync([result_key])
           result_for(result_key)
         end
@@ -68,7 +80,7 @@ def load_all(values)
           k = result_key_for(v)
           result_keys << k
           if !@results.key?(k)
-            @pending[k] ||= v
+            @pending[k] ||= normalize_fetch_key(v)
             pending_keys << k
           end
         }

spec/graphql/dataloader/active_record_association_source_spec.rb

@@ -105,5 +105,26 @@
       vulfpeck = d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :thing).load(wilco)
       assert_equal ::Band.find(1), vulfpeck
     end
+
+    if Rails::VERSION::STRING > "7.1" # not supported in <7.1
+      it_dataloads "loads with composite primary keys and warms the cache" do |d|
+        my_first_car = ::Album.find(2)
+        homey = ::Album.find(4)
+        log = with_active_record_log(colorize: false) do
+          vulfpeck, chon = d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :composite_band).load_all([my_first_car, homey])
+          assert_equal "Vulfpeck", vulfpeck.name
+          assert_equal "Chon", chon.name
+        end
+
+        assert_includes log, '[["name", "Vulfpeck"], ["name", "Chon"], ["genre", 0]]'
+
+
+        log = with_active_record_log(colorize: false) do
+          d.with(GraphQL::Dataloader::ActiveRecordSource, CompositeBand).load_all([["Vulfpeck", "rock"], ["Chon", :rock]])
+        end
+
+        assert_equal "", log
+      end
+    end
   end
 end

spec/graphql/dataloader/active_record_source_spec.rb

@@ -64,6 +64,23 @@
         assert_includes log, 'SELECT "bands".* FROM "bands" WHERE "bands"."name" = ?  [["name", "Vulfpeck"]]'
       end
 
+      if Rails::VERSION::STRING > "7.1" # not supported in <7.1
+        it_dataloads "uses composite primary keys" do |d|
+          log = with_active_record_log(colorize: false) do
+            r1 = d.with(GraphQL::Dataloader::ActiveRecordSource, CompositeBand).load(["Chon", :rock])
+            assert_equal "Chon", r1.name
+            assert_equal ["Chon", "rock"], r1.id
+            if Rails::VERSION::STRING > "8"
+              assert_equal 3, r1["id"]
+            else
+              assert_equal 3, r1._read_attribute("id")
+            end
+          end
+
+          assert_includes log, 'SELECT "bands".* FROM "bands" WHERE "bands"."name" = ? AND "bands"."genre" = ?  [["name", "Chon"], ["genre", 0]]'
+        end
+      end
+
       it_dataloads "uses specified find_by columns" do |d|
         log = with_active_record_log(colorize: false) do
           r1 = d.with(GraphQL::Dataloader::ActiveRecordSource, Band, find_by: :name).load("Chon")
@@ -99,11 +116,5 @@
         assert_equal "", log
       end
     end
-
-    describe "in queries" do
-      it "loads records with dataload_record"
-
-      it "accepts custom find-by with dataload_record"
-    end
   end
 end

spec/graphql/tracing/snapshots/example-rails-7-0.json

@@ -1487,13 +1487,13 @@
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "object",
-            "stringValue": "id"
+            "stringValue": "@type_for_column"
           },
           {
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "result",
-            "stringValue": "resolved_type"
+            "stringValue": "id"
           }
         ],
         "type": "TYPE_SLICE_BEGIN",
@@ -1556,13 +1556,13 @@
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "object",
-            "stringValue": "id"
+            "stringValue": "@type_for_column"
           },
           {
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "result",
-            "stringValue": null
+            "stringValue": "resolved_type"
           }
         ],
         "type": "TYPE_SLICE_BEGIN",
@@ -2961,6 +2961,11 @@
             "name": "@find_by",
             "stringValue": null
           },
+          {
+            "nameIid": "10101010101010",
+            "boolValue": false,
+            "name": "@find_by_many"
+          },
           {
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
@@ -3014,6 +3019,10 @@
             "iid": "10101010101010",
             "name": "@find_by"
           },
+          {
+            "iid": "10101010101010",
+            "name": "@find_by_many"
+          },
           {
             "iid": "10101010101010",
             "name": "@type_for_column"
@@ -3610,6 +3619,11 @@
             "name": "@find_by",
             "stringValue": null
           },
+          {
+            "nameIid": "10101010101010",
+            "boolValue": false,
+            "name": "@find_by_many"
+          },
           {
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
@@ -5768,6 +5782,11 @@
             "name": "@find_by",
             "stringValue": null
           },
+          {
+            "nameIid": "10101010101010",
+            "boolValue": false,
+            "name": "@find_by_many"
+          },
           {
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
@@ -6876,7 +6895,7 @@
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "result",
-            "stringValue": "id"
+            "stringValue": "@type_for_column"
           }
         ],
         "type": "TYPE_SLICE_BEGIN",
@@ -7845,6 +7864,11 @@
             "name": "@find_by",
             "stringValue": null
           },
+          {
+            "nameIid": "10101010101010",
+            "boolValue": false,
+            "name": "@find_by_many"
+          },
           {
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
@@ -8291,7 +8315,7 @@
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "result",
-            "stringValue": "id"
+            "stringValue": "@type_for_column"
           }
         ],
         "type": "TYPE_SLICE_BEGIN",
@@ -10573,13 +10597,13 @@
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "object",
-            "stringValue": "id"
+            "stringValue": "@type_for_column"
           },
           {
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "result",
-            "stringValue": "resolved_type"
+            "stringValue": "id"
           }
         ],
         "type": "TYPE_SLICE_BEGIN",
@@ -11035,13 +11059,13 @@
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "object",
-            "stringValue": "id"
+            "stringValue": "@type_for_column"
           },
           {
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",
             "name": "result",
-            "stringValue": "resolved_type"
+            "stringValue": "id"
           }
         ],
         "type": "TYPE_SLICE_BEGIN",
@@ -11660,6 +11684,11 @@
             "name": "@find_by",
             "stringValue": null
           },
+          {
+            "nameIid": "10101010101010",
+            "boolValue": false,
+            "name": "@find_by_many"
+          },
           {
             "nameIid": "10101010101010",
             "stringValueIid": "10101010101010",