Merge pull request #3106 from alphagov/towers/batch-graphql-edition-q…

…ueries

Batch link expansion queries

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}')" }.join(",")
+      edition_id_tuples = editions_and_link_types.map { |edition, link_type| "(#{edition.id},'#{link_type}')" }.join(",")
 
-      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))
+        .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))
+        .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
 

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
 

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,9 +49,32 @@
     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
   end
+
+  it "returns links from only the requested content store" do
+    target_edition_1 = create(:edition, content_store: "live")
+    target_edition_2 = create(:edition, content_store: "live")
+    target_edition_3 = create(:edition, content_store: "draft")
+    target_edition_4 = create(:edition, content_store: "draft")
+
+    source_edition = create(:edition,
+                            content_store: "draft",
+                            links_hash: {
+                              "test_link" => [target_edition_1.content_id, target_edition_3.content_id],
+                            })
+
+    link_set = create(:link_set, content_id: source_edition.content_id)
+    create(:link, link_set:, target_content_id: target_edition_2.content_id, link_type: "test_link")
+    create(:link, 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, content_store: source_edition.content_store).request([source_edition, "test_link"])
+
+      expect(request.load).to eq([target_edition_3, target_edition_4])
+    end
+  end
 end