Improve performance of stages index migration on PostgreSQL

1 files changed, 28 insertions, 9 deletions

diff --git a/lib/gitlab/background_migration/migrate_stage_index.rb b/lib/gitlab/background_migration/migrate_stage_index.rb
index c5f0fb5e8f8..6439ab3e42a 100644
--- a/lib/gitlab/background_migration/migrate_stage_index.rb
+++ b/lib/gitlab/background_migration/migrate_stage_index.rb
@@ -11,15 +11,34 @@ module Gitlab
       end
 
       def perform(start_id, stop_id)
-        sql = <<~SQL
-          UPDATE ci_stages
-            SET index =
-              (SELECT stage_idx FROM ci_builds
-                WHERE ci_builds.stage_id = ci_stages.id
-                GROUP BY ci_builds.stage_idx ORDER BY COUNT(*) DESC LIMIT 1)
-          WHERE ci_stages.id BETWEEN #{start_id.to_i} AND #{stop_id.to_i}
-            AND ci_stages.index IS NULL
-        SQL
+        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 =
+                (SELECT stage_idx FROM ci_builds
+                  WHERE ci_builds.stage_id = ci_stages.id
+                  GROUP BY ci_builds.stage_idx ORDER BY COUNT(*) DESC LIMIT 1)
+            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