Another heathy fixfix-retried-for-mysql-v2

2 files changed, 29 insertions, 31 deletions

diff --git a/db/post_migrate/20170503004427_upate_retried_for_ci_build.rb b/db/post_migrate/20170503004427_upate_retried_for_ci_build.rb
index 738e46b9207..a0eb4fa818c 100644
--- a/db/post_migrate/20170503004427_upate_retried_for_ci_build.rb
+++ b/db/post_migrate/20170503004427_upate_retried_for_ci_build.rb
@@ -8,10 +8,10 @@ class UpateRetriedForCiBuild < ActiveRecord::Migration
   def up
     disable_statement_timeout
 
-    if Gitlab::Database.mysql?
-      up_mysql
-    else
-      up_postgres
+    with_temporary_partial_index do
+      update_column_in_batches(:ci_builds, :retried, &method(:retried_query)) do |table, query|
+        query = query.where(table[:retried].eq(nil))
+      end
     end
   end
 
@@ -20,36 +20,28 @@ class UpateRetriedForCiBuild < ActiveRecord::Migration
 
   private
 
-  def up_mysql
-    # This is a trick to overcome MySQL limitation: 
-    # Mysql2::Error: Table 'ci_builds' is specified twice, both as a target for 'UPDATE' and as a separate source for data
-    # However, this leads to create a temporary table from `max(ci_builds.id)` which is slow and do full database update
-    execute <<-SQL.strip_heredoc
-      UPDATE ci_builds SET retried=
-        (id NOT IN (
-          SELECT * FROM (SELECT MAX(ci_builds.id) FROM ci_builds GROUP BY commit_id, name) AS latest_jobs
-        ))
-      WHERE retried IS NULL
+  def retried_query(table, query, start_id, stop_id)
+    latest_commits = <<-SQL.strip_heredoc
+      SELECT DISTINCT ci_builds3.commit_id
+        FROM ci_builds3
+        WHERE #{start_id} <= id
     SQL
-  end
 
-  def up_postgres
-    with_temporary_partial_index do
-      latest_id = <<-SQL.strip_heredoc
-        SELECT MAX(ci_builds2.id)
-          FROM ci_builds ci_builds2
-          WHERE ci_builds.commit_id=ci_builds2.commit_id
-            AND ci_builds.name=ci_builds2.name
-      SQL
+    latest_commits = "#{latest_commits} AND id < #{stop_id}" if stop_id
 
-      # This is slow update as it does single-row query
-      # This is designed to be run as idle, or a post deployment migration
-      is_retried = Arel.sql("((#{latest_id}) != ci_builds.id)")
+    latest_ids = <<-SQL.strip_heredoc
+      SELECT MAX(ci_builds2.id) FROM ci_builds2
+        WHERE commit_id IN (#{latest_commits})
+        GROUP BY commit_id, name
+    SQL
 
-      update_column_in_batches(:ci_builds, :retried, is_retried) do |table, query|
-        query.where(table[:retried].eq(nil))
-      end
-    end
+    latest_ids = <<-SQL.strip_heredoc
+      SELECT * FROM (#{latest_ids}) AS latest_ids
+    SQL
+
+    value = Arel.sql("(ci_builds.id NOT IN (#{latest_ids}))")
+    
+    query.set([[table[:retried], value]])
   end
 
   def with_temporary_partial_index
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index e76c9abbe04..199cbd71f5d 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -184,7 +184,6 @@ module Gitlab
 
           update_arel = Arel::UpdateManager.new(ActiveRecord::Base).
             table(table).
-            set([[table[column], value]]).
             where(table[:id].gteq(start_id))
 
           if stop_row
@@ -195,6 +194,13 @@ module Gitlab
 
           update_arel = yield table, update_arel if block_given?
 
+          update_arel = 
+            if value.is_a?(Proc)
+              value.call(table, update_arel, start_id, stop_id)
+            else
+              update_arel.set([[table[column], value]])
+            end
+
           execute(update_arel.to_sql)
 
           # There are no more rows left to update.