Reduce dataloader queries

app/graphql/sources/linked_to_editions_source.rb

@@ -7,36 +7,79 @@ def initialize(content_store:)
     # rubocop:enable Lint/MissingSuper
 
     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_selections = {
+        links: %i[link_type position],
+        documents: %i[content_id],
+      }
+      edition_id_tuples = []
+      content_id_tuples = []
+      link_types_map = {}
 
-      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)
+      editions_and_link_types.each do |edition, link_type|
+        edition_id_tuples.push("(#{edition.id},'#{link_type}')")
+        content_id_tuples.push("('#{edition.content_id}','#{link_type}')")
+        link_types_map[[edition.content_id, link_type]] = []
+      end
 
-      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_set_links_target_editions = Edition
+        .joins(document: { reverse_links: :link_set })
+        .where(
+          '("link_sets"."content_id", "links"."link_type") IN (?)',
+          Arel.sql(content_id_tuples.join(",")),
+        )
+        .where(
+          editions: { content_store: @content_store },
+          documents: { locale: "en" },
+        )
+        .select(
+          "editions.*",
+          all_selections,
+          { link_sets: { content_id: :source_content_id } },
+        )
 
-      all_links = edition_links + link_set_links
+      edition_links_target_editions = Edition
+        .joins(document: :reverse_links)
+        .joins(
+          <<~SQL,
+            INNER JOIN editions source_editions
+            ON source_editions.id = links.edition_id
+          SQL
+        )
+        .joins(
+          <<~SQL,
+            INNER JOIN documents source_documents
+            ON source_documents.id = source_editions.document_id
+          SQL
+        )
+        .where(
+          '("source_editions"."id", "links"."link_type") IN (?)',
+          Arel.sql(edition_id_tuples.join(",")),
+        )
+        .where(
+          editions: { content_store: @content_store },
+          documents: { locale: "en" },
+        )
+        .select(
+          "editions.*",
+          all_selections,
+          { source_documents: { content_id: :source_content_id } },
+        )
 
-      link_types_map = editions_and_link_types.map { [_1.content_id, _2] }.index_with { [] }
+      all_editions = Edition
+        .from(
+          <<~SQL,
+            (
+              #{link_set_links_target_editions.to_sql}
+              UNION
+              #{edition_links_target_editions.to_sql}
+            ) AS editions
+          SQL
+        )
+        .order(link_type: :asc, position: :asc)
 
-      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))
+      all_editions.each_with_object(link_types_map) { |edition, hash|
+        hash[[edition.source_content_id, edition.link_type]] << edition
       }.values
     end
-
-  private
-
-    def editions_for_link(link)
-      link.target_documents.map { |document| document.send(@content_store) }
-    end
   end
 end

app/graphql/sources/reverse_linked_to_editions_source.rb

@@ -7,27 +7,51 @@ def initialize(content_store:)
     # rubocop:enable Lint/MissingSuper
 
     def fetch(editions_and_link_types)
-      content_id_tuples = editions_and_link_types.map { |edition, link_type| "('#{edition.content_id}','#{link_type}')" }.join(",")
+      all_selections = {
+        links: %i[target_content_id link_type link_set_id edition_id],
+        documents: %i[content_id],
+      }
+      content_id_tuples = []
+      link_types_map = {}
 
-      all_links = Link
-        .where('("links"."target_content_id", "links"."link_type") IN (?)', Arel.sql(content_id_tuples))
-        .includes(source_documents: @content_store)
+      editions_and_link_types.each do |edition, link_type|
+        content_id_tuples.push("('#{edition.content_id}','#{link_type}')")
+        link_types_map[[edition.content_id, link_type]] = []
+      end
 
-      link_types_map = editions_and_link_types.map { [_1.content_id, _2] }.index_with { [] }
+      link_set_links_source_editions = Edition
+        .joins(:document)
+        .joins("INNER JOIN link_sets ON link_sets.content_id = documents.content_id")
+        .joins("INNER JOIN links ON links.link_set_id = link_sets.id")
+        .where(editions: { content_store: @content_store })
+        .where(
+          '("links"."target_content_id", "links"."link_type") IN (?)',
+          Arel.sql(content_id_tuples.join(",")),
+        )
+        .select("editions.*", all_selections)
 
-      all_links.each_with_object(link_types_map) { |link, hash|
-        if link.link_set
-          hash[[link.target_content_id, link.link_type]].concat(editions_for_link_set_link(link))
-        elsif link.edition
-          hash[[link.target_content_id, link.link_type]] << link.edition
-        end
-      }.values
-    end
+      edition_links_source_editions = Edition
+        .joins(:document, :links)
+        .where(editions: { content_store: @content_store })
+        .where(
+          '("links"."target_content_id", "links"."link_type") IN (?)',
+          Arel.sql(content_id_tuples.join(",")),
+        )
+        .select("editions.*", all_selections)
 
-  private
+      all_editions = Edition.from(
+        <<~SQL,
+          (
+            #{link_set_links_source_editions.to_sql}
+            UNION
+            #{edition_links_source_editions.to_sql}
+          ) AS editions
+        SQL
+      )
 
-    def editions_for_link_set_link(link)
-      link.source_documents.map { |document| document.send(@content_store) }
+      all_editions.each_with_object(link_types_map) { |edition, hash|
+        hash[[edition.target_content_id, edition.link_type]] << edition
+      }.values
     end
   end
 end

app/models/edition.rb

@@ -75,7 +75,15 @@ class Edition < ApplicationRecord
   validates_with StateForDocumentValidator
   validates_with RoutesAndRedirectsValidator
 
-  delegate :content_id, :locale, to: :document
+  def self.attribute_or_delegate(*names, to:)
+    names.each do |name|
+      define_method(name) do
+        attribute(name.to_s) || method(to).call.method(name).call
+      end
+    end
+  end
+
+  attribute_or_delegate :content_id, :locale, to: :document
 
   def auth_bypass_ids_are_uuids
     unless auth_bypass_ids.all? { |id| UuidValidator.valid?(id) }

spec/factories/link.rb

@@ -1,6 +1,6 @@
 FactoryBot.define do
   factory :link do
-    link_set
+    link_set          { create(:link_set) unless edition }
     target_content_id { SecureRandom.uuid }
     link_type         { "organisations" }
   end

spec/graphql/sources/linked_to_editions_source_spec.rb

@@ -12,7 +12,7 @@
     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_1, target_edition_3])
+      expect(request.load).to match_array([target_edition_1, target_edition_3])
     end
   end
 
@@ -30,7 +30,7 @@
     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_1, target_edition_3])
+      expect(request.load).to match_array([target_edition_1, target_edition_3])
     end
   end
 
@@ -51,7 +51,7 @@
     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_1, target_edition_3])
+      expect(request.load).to match_array([target_edition_1, target_edition_3])
     end
   end
 
@@ -74,7 +74,35 @@
     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])
+      expect(request.load).to match_array([target_edition_3, target_edition_4])
+    end
+  end
+
+  it "returns editions in order of their associated link's `position`" do
+    target_edition_0 = create(:edition)
+    target_edition_1 = create(:edition)
+    target_edition_2 = create(:edition)
+    target_edition_3 = create(:edition)
+
+    source_edition = create(:edition, content_store: "draft")
+    create(:link, edition: source_edition, target_content_id: target_edition_1.content_id, position: 1, link_type: "test_link")
+    create(:link, edition: source_edition, target_content_id: target_edition_3.content_id, position: 3, link_type: "test_link")
+
+    link_set = create(:link_set, content_id: source_edition.content_id)
+    create(:link, link_set:, target_content_id: target_edition_0.content_id, position: 0, link_type: "test_link")
+    create(:link, link_set:, target_content_id: target_edition_2.content_id, position: 2, 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",
+        %i[id base_path title document_id],
+      ])
+
+      expect(request.load).to eq([target_edition_0, target_edition_1, target_edition_2, target_edition_3])
     end
   end
 end