Clean up batch loading by introducing BatchKey

app/graphql/types/ci/stage_type.rb

@@ -20,26 +20,20 @@ module Types
 
       # Issues one query per pipeline
       def groups(lookahead:)
-        needs_selected = %i[nodes jobs nodes]
-          .reduce(lookahead) { |q, f| q.selection(f) }
-          .selects?(:needs)
-        key = [object.pipeline, object, needs_selected]
+        key = ::Gitlab::Graphql::BatchKey.new(object, lookahead, object_name: :stage)
 
         BatchLoader::GraphQL.for(key).batch(default_value: []) do |keys, loader|
-          by_pipeline = keys.group_by(&:first)
-          include_needs = keys.any? { |k| k[2] }
+          by_pipeline = keys.group_by(&:pipeline)
+          include_needs = keys.any? { |k| k.requires?(%i[nodes jobs nodes needs]) }
 
           by_pipeline.each do |pl, key_group|
             project = pl.project
-            stages = key_group.map(&:second).uniq
-            indexed = stages.index_by(&:id)
-
-            jobs_for_pipeline(pl, stages.map(&:id), include_needs).each do |stage_id, statuses|
-              stage = indexed[stage_id]
-              groups = ::Ci::Group.fabricate(project, stage, statuses)
-              # we don't know (and do not care) whether this set of jobs was
-              # loaded with needs preloaded as part of the key.
-              [true, false].each { |b| loader.call([pl, stage, b], groups) }
+            indexed = key_group.index_by(&:id)
+
+            jobs_for_pipeline(pl, indexed.keys, include_needs).each do |stage_id, statuses|
+              key = indexed[stage_id]
+              groups = ::Ci::Group.fabricate(project, key.stage, statuses)
+              loader.call(key, groups)
             end
           end
         end

lib/gitlab/graphql/batch_key.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    class BatchKey
+      attr_reader :object
+      delegate :hash, to: :object
+
+      def initialize(object, lookahead = nil, object_name: nil)
+        @object = object
+        @lookahead = lookahead
+        @object_name = object_name
+      end
+
+      def requires?(path)
+        return false unless @lookahead
+        return false unless path.present?
+
+        field = path.pop
+
+        path
+          .reduce(@lookahead) { |q, f| q.selection(f) }
+          .selects?(field)
+      end
+
+      def eql?(other)
+        other.is_a?(self.class) && object == other.object
+      end
+      alias_method :==, :eql?
+
+      def method_missing(method_name, *args, **kwargs)
+        return @object if method_name.to_sym == @object_name
+        return @object.public_send(method_name) if args.empty? && kwargs.empty? # rubocop: disable GitlabSecurity/PublicSend
+
+        super
+      end
+    end
+  end
+end

spec/lib/gitlab/graphql/batch_key_spec.rb

+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+require 'test_prof/recipes/rspec/let_it_be'
+
+RSpec.describe ::Gitlab::Graphql::BatchKey do
+  let_it_be(:rect) { Struct.new(:len, :width) }
+  let_it_be(:circle) { Struct.new(:radius) }
+  let(:lookahead) { nil }
+  let(:object) { rect.new(2, 3) }
+
+  subject { described_class.new(object, lookahead, object_name: :rect) }
+
+  it 'is equal to keys of the same object, regardless of lookahead or object name' do
+    expect(subject).to eq(described_class.new(rect.new(2, 3)))
+    expect(subject).to eq(described_class.new(rect.new(2, 3), :anything))
+    expect(subject).to eq(described_class.new(rect.new(2, 3), lookahead, object_name: :does_not_matter))
+    expect(subject).not_to eq(described_class.new(rect.new(2, 4)))
+    expect(subject).not_to eq(described_class.new(circle.new(10)))
+  end
+
+  it 'delegates attribute lookup methods to the inner object' do
+    other = rect.new(2, 3)
+
+    expect(subject.hash).to eq(other.hash)
+    expect(subject.len).to eq(other.len)
+    expect(subject.width).to eq(other.width)
+  end
+
+  it 'allows the object to be named more meaningfully' do
+    expect(subject.object).to eq(object)
+    expect(subject.object).to eq(subject.rect)
+  end
+
+  it 'works as a hash key' do
+    h = { subject => :foo }
+
+    expect(h[described_class.new(object)]).to eq(:foo)
+  end
+
+  describe '#requires?' do
+    it 'returns false if the lookahead was not provided' do
+      expect(subject.requires?([:foo])).to be(false)
+    end
+
+    context 'lookahead was provided' do
+      let(:lookahead) { double(:Lookahead) }
+
+      before do
+        allow(lookahead).to receive(:selection).with(Symbol).and_return(lookahead)
+      end
+
+      it 'returns false if the path is empty' do
+        expect(subject.requires?([])).to be(false)
+      end
+
+      context 'it selects the field' do
+        before do
+          allow(lookahead).to receive(:selects?).with(Symbol).once.and_return(true)
+        end
+
+        it 'returns true' do
+          expect(subject.requires?(%i[foo bar baz])).to be(true)
+        end
+      end
+
+      context 'it does not select the field' do
+        before do
+          allow(lookahead).to receive(:selects?).with(Symbol).once.and_return(false)
+        end
+
+        it 'returns false' do
+          expect(subject.requires?(%i[foo bar baz])).to be(false)
+        end
+      end
+    end
+  end
+end