Reduce query counts in GraphQL dataloader sources

We found that the previous eager loading approach wasn't using joins
effectively, so we're more manually crafting the queries to join in an
efficient way. We also found that some associations we're still lazy
loaded after the initial queries - this new approach avoids that issue

For now we're selecting all columns from the editions table. We'll
reduce the selections down to a subset in a later change

The change to using a UNION affects the order of the resulting editions,
so we've added a test to document our expectations

Co-authored-by: Mike Patrick <mike.patrick2@digital.cabinet-office.gov.uk>

Author: yndajas

app/graphql/sources/linked_to_editions_source.rb

@@ -7,6 +7,10 @@ def initialize(content_store:)
     # rubocop:enable Lint/MissingSuper
 
     def fetch(editions_and_link_types)
+      all_selections = {
+        links: %i[link_type position],
+        documents: %i[content_id],
+      }
       edition_id_tuples = []
       content_id_tuples = []
       link_types_map = {}
@@ -17,37 +21,65 @@ def fetch(editions_and_link_types)
         link_types_map[[edition.content_id, link_type]] = []
       end
 
-      edition_links = Link
-        .joins(:edition, { target_documents: @content_store })
-        .includes(:edition, { target_documents: @content_store })
+      link_set_links_target_editions = Edition
+        .joins(document: { reverse_links: :link_set })
         .where(
-          '("editions"."id", "links"."link_type") IN (?)',
-          Arel.sql(edition_id_tuples.join(",")),
+          '("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 } },
         )
-        .where(target_documents: { locale: "en" })
-        .order(link_type: :asc, position: :asc)
 
-      link_set_links = Link
-        .joins(:link_set, { target_documents: @content_store })
-        .includes(:link_set, { target_documents: @content_store })
+      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(
-          '("link_sets"."content_id", "links"."link_type") IN (?)',
-          Arel.sql(content_id_tuples.join(",")),
+          '("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 } },
         )
-        .where(target_documents: { locale: "en" })
-        .order(link_type: :asc, position: :asc)
 
-      all_links = edition_links + link_set_links
+      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,6 +7,10 @@ def initialize(content_store:)
     # rubocop:enable Lint/MissingSuper
 
     def fetch(editions_and_link_types)
+      all_selections = {
+        links: %i[target_content_id link_type link_set_id edition_id],
+        documents: %i[content_id],
+      }
       content_id_tuples = []
       link_types_map = {}
 
@@ -15,26 +19,39 @@ def fetch(editions_and_link_types)
         link_types_map[[edition.content_id, link_type]] = []
       end
 
-      all_links = Link
+      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(",")),
         )
-        .includes(source_documents: @content_store)
+        .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

spec/graphql/sources/linked_to_editions_source_spec.rb

@@ -77,4 +77,32 @@
       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