Use database query to calculate average stage position

2 files changed, 26 insertions, 10 deletions

diff --git a/app/models/ci/stage.rb b/app/models/ci/stage.rb
index a80c72d3452..5a77a909b9d 100644
--- a/app/models/ci/stage.rb
+++ b/app/models/ci/stage.rb
@@ -24,12 +24,14 @@ module Ci
       self.status = DEFAULT_STATUS if self.status.nil?
     end
 
-    before_validation do
-      next unless index.nil?
+    before_validation unless: :importing? do
+      next if index.present?
 
-      statuses.pluck(:stage_idx).tap do |indices|
-        self.index = indices.max_by { |index| indices.count(index) }
-      end
+      self.index = statuses.select(:stage_idx)
+        .where('stage_idx IS NOT NULL')
+        .group(:stage_idx)
+        .order('COUNT(*) DESC')
+        .first&.stage_idx.to_i
     end
 
     state_machine :status, initial: :created do
diff --git a/spec/models/ci/stage_spec.rb b/spec/models/ci/stage_spec.rb
index ae5ab6fe2ab..20db6394e51 100644
--- a/spec/models/ci/stage_spec.rb
+++ b/spec/models/ci/stage_spec.rb
@@ -92,16 +92,30 @@ describe Ci::Stage, :models do
     context 'when stage has been imported and does not have index set' do
       before do
         stage.update_column(:index, nil)
+      end
+
+      context 'when stage has statuses' do
+        before do
+          create(:ci_build, :running, stage_id: stage.id, stage_idx: 10)
+        end
+
+        it 'recalculates index before updating status' do
+          expect(stage.reload.index).to be_nil
 
-        create(:ci_build, :running, stage_id: stage.id, stage_idx: 10)
+          stage.update_status
+
+          expect(stage.reload.index).to eq 10
+        end
       end
 
-      it 'recalculates index before updating status' do
-        expect(stage.reload.index).to be_nil
+      context 'when stage does not have statuses' do
+        it 'fallbacks to zero' do
+          expect(stage.reload.index).to be_nil
 
-        stage.update_status
+          stage.update_status
 
-        expect(stage.reload.index).to eq 10
+          expect(stage.reload.index).to eq 0
+        end
       end
     end
   end