Implement LazyRelationLoader abstraction for GraphQL

This change introduces a new loader abstraction to be able to preload
`has_many` and `has_one` associations of records lazily.

lib/gitlab/graphql/loaders/lazy_relation_loader.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module Loaders
+      class LazyRelationLoader
+        class << self
+          attr_accessor :model, :association
+
+          # Automatically register the inheriting
+          # classes to GitlabSchema as lazy objects.
+          def inherited(klass)
+            GitlabSchema.lazy_resolve(klass, :load)
+          end
+        end
+
+        def initialize(query_ctx, object, **kwargs)
+          @query_ctx = query_ctx
+          @object = object
+          @kwargs = kwargs
+
+          query_ctx[loader_cache_key] ||= Registry.new(relation(**kwargs))
+          query_ctx[loader_cache_key].register(object)
+        end
+
+        # Returns an instance of `RelationProxy` for the object (parent model).
+        # The returned object behaves like an Active Record relation to support
+        # keyset pagination.
+        def load
+          case reflection.macro
+          when :has_many
+            relation_proxy
+          when :has_one
+            relation_proxy.last
+          else
+            raise 'Not supported association type!'
+          end
+        end
+
+        private
+
+        attr_reader :query_ctx, :object, :kwargs
+
+        delegate :model, :association, to: :"self.class"
+
+        # Implement this one if you want to filter the relation
+        def relation(**)
+          base_relation
+        end
+
+        def loader_cache_key
+          @loader_cache_key ||= self.class.name.to_s + kwargs.sort.to_s
+        end
+
+        def base_relation
+          placeholder_record.association(association).scope
+        end
+
+        # This will only work for HasMany and HasOne associations for now
+        def placeholder_record
+          model.new(reflection.active_record_primary_key => 0)
+        end
+
+        def reflection
+          model.reflections[association.to_s]
+        end
+
+        def relation_proxy
+          RelationProxy.new(object, query_ctx[loader_cache_key])
+        end
+      end
+    end
+  end
+end

lib/gitlab/graphql/loaders/lazy_relation_loader/registry.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module Loaders
+      class LazyRelationLoader
+        class Registry
+          PrematureQueryExecutionTriggered = Class.new(RuntimeError)
+          # Following methods are Active Record kicker methods which fire SQL query.
+          # We can support some of them with TopNLoader but for now restricting their use
+          # as we don't have a use case.
+          PROHIBITED_METHODS = (
+            ActiveRecord::FinderMethods.instance_methods(false) +
+            ActiveRecord::Calculations.instance_methods(false)
+          ).to_set.freeze
+
+          def initialize(relation)
+            @parents = []
+            @relation = relation
+            @records = []
+            @loaded = false
+          end
+
+          def register(object)
+            @parents << object
+          end
+
+          def method_missing(method_name, ...)
+            raise PrematureQueryExecutionTriggered if PROHIBITED_METHODS.include?(method_name)
+
+            result = relation.public_send(method_name, ...) # rubocop:disable GitlabSecurity/PublicSend
+
+            if result.is_a?(ActiveRecord::Relation) # Spawn methods generate a new relation (e.g. where, limit)
+              @relation = result
+
+              return self
+            end
+
+            result
+          end
+
+          def respond_to_missing?(method_name, include_private = false)
+            relation.respond_to?(method_name, include_private)
+          end
+
+          def load
+            return records if loaded
+
+            @loaded = true
+            @records = TopNLoader.load(relation, parents)
+          end
+
+          def for(object)
+            load.select { |record| record[foreign_key] == object[active_record_primary_key] }
+                .tap { |records| set_inverse_of(object, records) }
+          end
+
+          private
+
+          attr_reader :parents, :relation, :records, :loaded
+
+          delegate :proxy_association, to: :relation, private: true
+          delegate :reflection, to: :proxy_association, private: true
+          delegate :active_record_primary_key, :foreign_key, to: :reflection, private: true
+
+          def set_inverse_of(object, records)
+            records.each do |record|
+              object.association(reflection.name).set_inverse_instance(record)
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/graphql/loaders/lazy_relation_loader/relation_proxy.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module Loaders
+      class LazyRelationLoader
+        # Proxies all the method calls to Registry instance.
+        # The main purpose of having this is that calling load
+        # on an instance of this class will only return the records
+        # associated with the main Active Record model.
+        class RelationProxy
+          def initialize(object, registry)
+            @object = object
+            @registry = registry
+          end
+
+          def load
+            registry.for(object)
+          end
+          alias_method :to_a, :load
+
+          def last(limit = 1)
+            result = registry.limit(limit)
+                           .reverse_order!
+                           .for(object)
+
+            return result.first if limit == 1 # This is the Active Record behavior
+
+            result
+          end
+
+          private
+
+          attr_reader :registry, :object
+
+          # Delegate everything to registry
+          def method_missing(method_name, ...)
+            result = registry.public_send(method_name, ...) # rubocop:disable GitlabSecurity/PublicSend
+
+            return self if result == registry
+
+            result
+          end
+
+          def respond_to_missing?(method_name, include_private = false)
+            registry.respond_to?(method_name, include_private)
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/graphql/loaders/lazy_relation_loader/top_n_loader.rb

+# frozen_string_literal: true
+
+# rubocop:disable CodeReuse/ActiveRecord
+module Gitlab
+  module Graphql
+    module Loaders
+      class LazyRelationLoader
+        # Loads the top-n records for each given parent record.
+        # For example; if you want to load only 5 confidential issues ordered by
+        # their updated_at column per project for a list of projects by issuing only a single
+        # SQL query then this class can help you.
+        # Note that the limit applies per parent record which means that if you apply limit as 5
+        # for 10 projects, this loader will load 50 records in total.
+        class TopNLoader
+          def self.load(original_relation, parents)
+            new(original_relation, parents).load
+          end
+
+          def initialize(original_relation, parents)
+            @original_relation = original_relation
+            @parents = parents
+          end
+
+          def load
+            klass.select(klass.arel_table[Arel.star])
+                 .from(from)
+                 .joins("JOIN LATERAL (#{lateral_relation.to_sql}) AS #{klass.arel_table.name} ON true")
+                 .includes(original_includes)
+                 .preload(original_preload)
+                 .eager_load(original_eager_load)
+                 .load
+          end
+
+          private
+
+          attr_reader :original_relation, :parents
+
+          delegate :proxy_association, to: :original_relation, private: true
+          delegate :reflection, to: :proxy_association, private: true
+          delegate :klass, :foreign_key, :active_record, :active_record_primary_key,
+            to: :reflection, private: true
+
+          # This only works for HasMany and HasOne.
+          def lateral_relation
+            original_relation
+              .unscope(where: foreign_key) # unscoping the where condition generated for the placeholder_record.
+              .where(klass.arel_table[foreign_key].eq(active_record.arel_table[active_record_primary_key]))
+          end
+
+          def from
+            grouping_arel_node.as("#{active_record.arel_table.name}(#{active_record.primary_key})")
+          end
+
+          def grouping_arel_node
+            Arel::Nodes::Grouping.new(id_list_arel_node)
+          end
+
+          def id_list_arel_node
+            parent_ids.map { |id| [id] }
+                      .then { |ids| Arel::Nodes::ValuesList.new(ids) }
+          end
+
+          def parent_ids
+            parents.pluck(active_record.primary_key)
+          end
+
+          def original_includes
+            original_relation.includes_values
+          end
+
+          def original_preload
+            original_relation.preload_values
+          end
+
+          def original_eager_load
+            original_relation.eager_load_values
+          end
+        end
+      end
+    end
+  end
+end
+# rubocop:enable CodeReuse/ActiveRecord

lib/gitlab/graphql/pagination/connections.rb

@@ -13,6 +13,10 @@ def self.use(schema)
             ActiveRecord::Relation,
             Gitlab::Graphql::Pagination::Keyset::Connection)
 
+          schema.connections.add(
+            Gitlab::Graphql::Loaders::LazyRelationLoader::RelationProxy,
+            Gitlab::Graphql::Pagination::Keyset::Connection)
+
           schema.connections.add(
             Gitlab::Graphql::ExternallyPaginatedArray,
             Gitlab::Graphql::Pagination::ExternallyPaginatedArrayConnection)

spec/lib/gitlab/graphql/loaders/lazy_relation_loader/registry_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Graphql::Loaders::LazyRelationLoader::Registry, feature_category: :vulnerability_management do
+  describe '#respond_to?' do
+    let(:relation) { Project.all }
+    let(:registry) { described_class.new(relation) }
+
+    subject { registry.respond_to?(method_name) }
+
+    context 'when the relation responds to given method' do
+      let(:method_name) { :sorted_by_updated_asc }
+
+      it { is_expected.to be_truthy }
+    end
+
+    context 'when the relation does not respond to given method' do
+      let(:method_name) { :this_method_does_not_exist }
+
+      it { is_expected.to be_falsey }
+    end
+  end
+end

spec/lib/gitlab/graphql/loaders/lazy_relation_loader/relation_proxy_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Graphql::Loaders::LazyRelationLoader::RelationProxy, feature_category: :vulnerability_management do
+  describe '#respond_to?' do
+    let(:object) { double }
+    let(:registry) { instance_double(Gitlab::Graphql::Loaders::LazyRelationLoader::Registry) }
+    let(:relation_proxy) { described_class.new(object, registry) }
+
+    subject { relation_proxy.respond_to?(:foo) }
+
+    before do
+      allow(registry).to receive(:respond_to?).with(:foo, false).and_return(responds_to?)
+    end
+
+    context 'when the registry responds to given method' do
+      let(:responds_to?) { true }
+
+      it { is_expected.to be_truthy }
+    end
+
+    context 'when the registry does not respond to given method' do
+      let(:responds_to?) { false }
+
+      it { is_expected.to be_falsey }
+    end
+  end
+end

spec/lib/gitlab/graphql/loaders/lazy_relation_loader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Graphql::Loaders::LazyRelationLoader, feature_category: :vulnerability_management do
+  let(:query_context) { {} }
+  let(:args) { {} }
+
+  let_it_be(:project) { create(:project) }
+
+  let(:loader) { loader_class.new(query_context, project, **args) }
+
+  describe '#load' do
+    subject(:load_relation) { loader.load }
+
+    context 'when the association is has many' do
+      let_it_be(:public_issue) { create(:issue, project: project) }
+      let_it_be(:confidential_issue) { create(:issue, :confidential, project: project) }
+
+      let(:loader_class) do
+        Class.new(described_class) do
+          self.model = Project
+          self.association = :issues
+
+          def relation(public_only: false)
+            relation = base_relation
+            relation = relation.public_only if public_only
+
+            relation
+          end
+        end
+      end
+
+      it { is_expected.to be_an_instance_of(described_class::RelationProxy) }
+
+      describe '#relation' do
+        subject { load_relation.load }
+
+        context 'without arguments' do
+          it { is_expected.to contain_exactly(public_issue, confidential_issue) }
+        end
+
+        context 'with arguments' do
+          let(:args) { { public_only: true } }
+
+          it { is_expected.to contain_exactly(public_issue) }
+        end
+      end
+
+      describe 'using the same context for different records' do
+        let_it_be(:another_project) { create(:project) }
+
+        let(:loader_for_another_project) { loader_class.new(query_context, another_project, **args) }
+        let(:records_for_another_project) { loader_for_another_project.load.load }
+        let(:records_for_project) { load_relation.load }
+
+        before do
+          loader # register the original loader to query context
+        end
+
+        it 'does not mix associated records' do
+          expect(records_for_another_project).to be_empty
+          expect(records_for_project).to contain_exactly(public_issue, confidential_issue)
+        end
+
+        it 'does not cause N+1 queries' do
+          expect { records_for_another_project }.not_to exceed_query_limit(1)
+        end
+      end
+
+      describe 'using Active Record querying methods' do
+        subject { load_relation.limit(1).load.count }
+
+        it { is_expected.to be(1) }
+      end
+
+      describe 'using Active Record finder methods' do
+        subject { load_relation.last(2) }
+
+        it { is_expected.to contain_exactly(public_issue, confidential_issue) }
+      end
+
+      describe 'calling a method that returns a non relation object' do
+        subject { load_relation.limit(1).limit_value }
+
+        it { is_expected.to be(1) }
+      end
+
+      describe 'calling a prohibited method' do
+        subject(:count) { load_relation.count }
+
+        it 'raises a `PrematureQueryExecutionTriggered` error' do
+          expect { count }.to raise_error(described_class::Registry::PrematureQueryExecutionTriggered)
+        end
+      end
+    end
+
+    context 'when the association is has one' do
+      let!(:project_setting) { create(:project_setting, project: project) }
+      let(:loader_class) do
+        Class.new(described_class) do
+          self.model = Project
+          self.association = :project_setting
+        end
+      end
+
+      it { is_expected.to eq(project_setting) }
+    end
+
+    context 'when the association is belongs to' do
+      let(:loader_class) do
+        Class.new(described_class) do
+          self.model = Project
+          self.association = :namespace
+        end
+      end
+
+      it 'raises error' do
+        expect { load_relation }.to raise_error(RuntimeError)
+      end
+    end
+  end
+end