Batch GraphQL link expansion queries

Previously LinkedToEditionsSource would make one query per edition (getting linked editions, given a list of link types). This change makes one query per parent edition, by getting the dataloader to batch queries across multiple editions and link types. For the prime minister page, the number of database queries reduces from 93 to 15, and the ActiveRecord execution time decreases from ~100ms to ~20ms, when run locally. Co-authored-by: Richard Towers <[email protected]>
alphagov · Jan 30, 2025 · 179143e · 179143e
1 parent 7e95514
commit 179143e
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 23 deletions.
diff --git a/app/graphql/sources/linked_to_editions_source.rb b/app/graphql/sources/linked_to_editions_source.rb
@@ -1,30 +1,35 @@
 module Sources
   class LinkedToEditionsSource < GraphQL::Dataloader::Source
     # rubocop:disable Lint/MissingSuper
-    def initialize(parent_object:)
-      @object = parent_object
-      @content_store = parent_object.content_store.to_sym
+    def initialize(content_store:)
+      @content_store = content_store.to_sym
     end
     # rubocop:enable Lint/MissingSuper
 
-    def fetch(link_types)
-      edition_links = @object.links.order(link_type: :asc, position: :asc)
+    def fetch(editions_and_link_types)
+      content_id_tuples = editions_and_link_types.map { |edition, link_type| "('#{edition.content_id}','#{link_type}')" }
+      edition_id_tuples = editions_and_link_types.map { |edition, link_type| "(#{edition.id},'#{link_type}')" }
 
-      all_links = if @object.document.link_set
-                    edition_links.or(@object.document.link_set.links)
-                  else
-                    edition_links
-                  end
+      edition_links = Link
+        .joins(:edition, { target_documents: @content_store })
+        .includes(:edition, { target_documents: @content_store })
+        .where('("editions"."id", "links"."link_type") IN (?)', Arel.sql(edition_id_tuples.join(",")))
+        .where(target_documents: { locale: "en" })
+        .order(link_type: :asc, position: :asc)
 
-      all_links_for_link_type_and_locale = all_links
-                                            .includes(target_documents: @content_store) # content_store is :live or :draft (a Document has_one Edition by that name)
-                                            .where(link_type: link_types)
-                                            .where(target_documents: { locale: "en" })
+      link_set_links = Link
+        .joins(:link_set, { target_documents: @content_store })
+        .includes(:link_set, { target_documents: @content_store })
+        .where('("link_sets"."content_id", "links"."link_type") IN (?)', Arel.sql(content_id_tuples.join(",")))
+        .where(target_documents: { locale: "en" })
+        .order(link_type: :asc, position: :asc)
 
-      link_types_map = link_types.index_with { [] }
+      all_links = edition_links + link_set_links
 
-      all_links_for_link_type_and_locale.each_with_object(link_types_map) { |link, hash|
-        hash[link.link_type].concat(editions_for_link(link))
+      link_types_map = editions_and_link_types.map { [_1.content_id, _2] }.index_with { [] }
+
+      all_links.each_with_object(link_types_map) { |link, hash|
+        hash[[(link.link_set || link.edition).content_id, link.link_type]].concat(editions_for_link(link))
       }.values
     end
 

diff --git a/app/graphql/types/base_object.rb b/app/graphql/types/base_object.rb
@@ -10,8 +10,8 @@ def self.links_field(field_name_and_link_type, graphql_field_type)
       field(field_name_and_link_type.to_sym, graphql_field_type)
 
       define_method(field_name_and_link_type.to_sym) do
-        dataloader.with(Sources::LinkedToEditionsSource, parent_object: object)
-          .load(field_name_and_link_type.to_s)
+        dataloader.with(Sources::LinkedToEditionsSource, content_store: object.content_store)
+          .load([object, field_name_and_link_type.to_s])
       end
     end
 

diff --git a/spec/graphql/sources/linked_to_editions_source_spec.rb b/spec/graphql/sources/linked_to_editions_source_spec.rb
@@ -10,7 +10,7 @@
     create(:link, link_set: link_set, target_content_id: target_edition_3.content_id, link_type: "test_link")
 
     GraphQL::Dataloader.with_dataloading do |dataloader|
-      request = dataloader.with(described_class, parent_object: source_edition).request("test_link")
+      request = dataloader.with(described_class, content_store: source_edition.content_store).request([source_edition, "test_link"])
 
       expect(request.load).to eq([target_edition_1, target_edition_3])
     end
@@ -28,7 +28,7 @@
                             })
 
     GraphQL::Dataloader.with_dataloading do |dataloader|
-      request = dataloader.with(described_class, parent_object: source_edition).request("test_link")
+      request = dataloader.with(described_class, content_store: source_edition.content_store).request([source_edition, "test_link"])
 
       expect(request.load).to eq([target_edition_1, target_edition_3])
     end
@@ -49,7 +49,7 @@
     create(:link, link_set: link_set, target_content_id: target_edition_3.content_id, link_type: "test_link")
 
     GraphQL::Dataloader.with_dataloading do |dataloader|
-      request = dataloader.with(described_class, parent_object: source_edition).request("test_link")
+      request = dataloader.with(described_class, content_store: source_edition.content_store).request([source_edition, "test_link"])
 
       expect(request.load).to eq([target_edition_1, target_edition_3])
     end
@@ -72,7 +72,7 @@
     create(:link, link_set: link_set, target_content_id: target_edition_4.content_id, link_type: "test_link")
 
     GraphQL::Dataloader.with_dataloading do |dataloader|
-      request = dataloader.with(described_class, parent_object: source_edition).request("test_link")
+      request = dataloader.with(described_class, content_store: source_edition.content_store).request([source_edition, "test_link"])
 
       expect(request.load).to eq([target_edition_3, target_edition_4])
     end