Fix N+1 query problem for previousStageJobsOrNeeds

There were three problems;
1. The relation name was previousStageJobsAndNeeds, then it was changed
to previousStageJobsOrNeeds. However, the include_needs condition
stayed the same.
2. Batch loader of Job.previous_stage_jobs was not working well.
3. N+1 tests were not working well

app/graphql/types/ci/job_type.rb

@@ -87,8 +87,10 @@ class JobType < BaseObject
                                                 description: 'Play path of the job.'
       field :playable, GraphQL::Types::Boolean, null: false, method: :playable?,
                                                 description: 'Indicates the job can be played.'
-      field :previous_stage_jobs_or_needs, Types::Ci::JobNeedUnion.connection_type, null: true,
-                                                                                    description: 'Jobs that must complete before the job runs. Returns `BuildNeed`, which is the needed jobs if the job uses the `needs` keyword, or the previous stage jobs otherwise.'
+      field :previous_stage_jobs_or_needs, Types::Ci::JobNeedUnion.connection_type,
+            null: true,
+            description: 'Jobs that must complete before the job runs. Returns `BuildNeed`, ' \
+                         'which is the needed jobs if the job uses the `needs` keyword, or the previous stage jobs otherwise.'
       field :ref_name, GraphQL::Types::String, null: true,
                                                description: 'Ref name of the job.'
       field :ref_path, GraphQL::Types::String, null: true,
@@ -179,7 +181,9 @@ def previous_stage_jobs
             stages = pipeline.stages.by_position(positions)
 
             stages.each do |stage|
-              loader.call([pipeline, stage.position], stage.latest_statuses)
+              # Without `.to_a`, the memoization will only preserve the activerecord relation object. And when there is
+              # a call, the SQL query will be executed again.
+              loader.call([pipeline, stage.position], stage.latest_statuses.to_a)
             end
           end
         end

app/graphql/types/ci/stage_type.rb

@@ -33,7 +33,7 @@ def groups(lookahead:)
           by_pipeline = keys.group_by(&:pipeline)
           include_needs = keys.any? do |k|
             k.requires?(%i[nodes jobs nodes needs]) ||
-            k.requires?(%i[nodes jobs nodes previousStageJobsAndNeeds])
+            k.requires?(%i[nodes jobs nodes previousStageJobsOrNeeds])
           end
 
           by_pipeline.each do |pl, key_group|

spec/requests/api/graphql/project/jobs_spec.rb

@@ -11,6 +11,9 @@
     create(:ci_pipeline, project: project, user: user)
   end
 
+  let!(:job1) { create(:ci_build, pipeline: pipeline, name: 'job 1') }
+  let!(:job2) { create(:ci_build, pipeline: pipeline, name: 'job 2') }
+
   let(:query) do
     <<~QUERY
     {
@@ -18,11 +21,6 @@
         jobs {
           nodes {
             name
-            previousStageJobsAndNeeds {
-              nodes {
-                name
-              }
-            }
           }
         }
       }
@@ -30,27 +28,10 @@
     QUERY
   end
 
-  it 'does not generate N+1 queries', :request_store, :use_sql_query_cache do
-    build_stage = create(:ci_stage, position: 1, name: 'build', project: project, pipeline: pipeline)
-    test_stage = create(:ci_stage, position: 2, name: 'test', project: project, pipeline: pipeline)
-    create(:ci_build, pipeline: pipeline, name: 'docker 1 2', ci_stage: build_stage)
-    create(:ci_build, pipeline: pipeline, name: 'docker 2 2', ci_stage: build_stage)
-    create(:ci_build, pipeline: pipeline, name: 'rspec 1 2', ci_stage: test_stage)
-    test_job = create(:ci_build, pipeline: pipeline, name: 'rspec 2 2', ci_stage: test_stage)
-    create(:ci_build_need, build: test_job, name: 'docker 1 2')
-
+  it 'fetches jobs' do
     post_graphql(query, current_user: user)
+    expect_graphql_errors_to_be_empty
 
-    control = ActiveRecord::QueryRecorder.new(skip_cached: false) do
-      post_graphql(query, current_user: user)
-    end
-
-    create(:ci_build, name: 'test-a', ci_stage: test_stage, pipeline: pipeline)
-    test_b_job = create(:ci_build, name: 'test-b', ci_stage: test_stage, pipeline: pipeline)
-    create(:ci_build_need, build: test_b_job, name: 'docker 2 2')
-
-    expect do
-      post_graphql(query, current_user: user)
-    end.not_to exceed_all_query_limit(control)
+    expect(graphql_data['project']['jobs']['nodes'].pluck('name')).to contain_exactly('job 1', 'job 2')
   end
 end

spec/requests/api/graphql/project/pipeline_spec.rb

@@ -337,16 +337,33 @@ def successful_pipeline
     end
   end
 
-  context 'N+1 queries on pipeline jobs' do
+  context 'N+1 queries on pipeline jobs.previousStageJobsOrNeeds' do
     let(:pipeline) { create(:ci_pipeline, project: project) }
 
     let(:fields) do
       <<~FIELDS
-      jobs {
+      stages {
         nodes {
-          previousStageJobsAndNeeds {
+          groups {
             nodes {
-              name
+              jobs {
+                nodes {
+                  previousStageJobsOrNeeds {
+                    nodes {
+                      ... on JobNeedUnion {
+                        ... on CiBuildNeed {
+                          id
+                          name
+                        }
+                        ... on CiJob {
+                          id
+                          name
+                        }
+                      }
+                    }
+                  }
+                }
+              }
             }
           }
         }
@@ -357,13 +374,15 @@ def successful_pipeline
     it 'does not generate N+1 queries', :request_store, :use_sql_query_cache do
       build_stage = create(:ci_stage, position: 1, name: 'build', project: project, pipeline: pipeline)
       test_stage = create(:ci_stage, position: 2, name: 'test', project: project, pipeline: pipeline)
-      create(:ci_build, pipeline: pipeline, name: 'docker 1 2', ci_stage: build_stage)
+
+      docker_1_2 = create(:ci_build, pipeline: pipeline, name: 'docker 1 2', ci_stage: build_stage)
       create(:ci_build, pipeline: pipeline, name: 'docker 2 2', ci_stage: build_stage)
       create(:ci_build, pipeline: pipeline, name: 'rspec 1 2', ci_stage: test_stage)
-      test_job = create(:ci_build, pipeline: pipeline, name: 'rspec 2 2', ci_stage: test_stage)
-      create(:ci_build_need, build: test_job, name: 'docker 1 2')
+      create(:ci_build, :dependent, needed: docker_1_2, pipeline: pipeline, name: 'rspec 2 2', ci_stage: test_stage)
 
+      # warm up
       post_graphql(query, current_user: current_user)
+      expect_graphql_errors_to_be_empty
 
       control = ActiveRecord::QueryRecorder.new(skip_cached: false) do
         post_graphql(query, current_user: current_user)
@@ -371,7 +390,7 @@ def successful_pipeline
 
       create(:ci_build, name: 'test-a', ci_stage: test_stage, pipeline: pipeline)
       test_b_job = create(:ci_build, name: 'test-b', ci_stage: test_stage, pipeline: pipeline)
-      create(:ci_build_need, build: test_b_job, name: 'docker 2 2')
+      create(:ci_build, :dependent, needed: test_b_job, pipeline: pipeline, name: 'docker 2 2', ci_stage: test_stage)
 
       expect do
         post_graphql(query, current_user: current_user)
@@ -424,6 +443,7 @@ def successful_pipeline
 
       # warm up
       post_graphql(query, current_user: current_user)
+      expect_graphql_errors_to_be_empty
 
       control = ActiveRecord::QueryRecorder.new(skip_cached: false) do
         post_graphql(query, current_user: current_user)