Improve performance of stages index migration on PostgreSQL

lib/gitlab/background_migration/migrate_stage_index.rb

@@ -11,6 +11,24 @@ module Gitlab
       end
 
       def perform(start_id, stop_id)
+        if Gitlab::Database.postgresql?
+          sql = <<~SQL
+            WITH freqs AS (
+              SELECT stage_id, stage_idx, COUNT(*) AS freq FROM ci_builds
+                WHERE stage_id BETWEEN #{start_id.to_i} AND #{stop_id.to_i}
+                  AND stage_idx IS NOT NULL
+                GROUP BY stage_id, stage_idx
+            ), indexes AS (
+              SELECT DISTINCT stage_id, last_value(stage_idx)
+                OVER (PARTITION BY stage_id ORDER BY freq ASC) AS index
+                FROM freqs
+            )
+
+            UPDATE ci_stages SET index = indexes.index
+              FROM indexes WHERE indexes.stage_id = ci_stages.id
+                AND ci_stages.index IS NULL;
+          SQL
+        else
           sql = <<~SQL
             UPDATE ci_stages
               SET index =
@@ -20,6 +38,7 @@ module Gitlab
             WHERE ci_stages.id BETWEEN #{start_id.to_i} AND #{stop_id.to_i}
               AND ci_stages.index IS NULL
           SQL
+        end
 
         ActiveRecord::Base.connection.execute(sql)
       end