Use a CTE to remove the query timeout

4 files changed, 113 insertions, 10 deletions

diff --git a/changelogs/unreleased/50359-activerecord-statementinvalid-pg-querycanceled-error-canceling-statement-due-to-statement-timeout.yml b/changelogs/unreleased/50359-activerecord-statementinvalid-pg-querycanceled-error-canceling-statement-due-to-statement-timeout.yml
new file mode 100644
index 00000000000..09ec4b8d73d
--- /dev/null
+++ b/changelogs/unreleased/50359-activerecord-statementinvalid-pg-querycanceled-error-canceling-statement-due-to-statement-timeout.yml
@@ -0,0 +1,5 @@
+---
+title: Fix timeout when running the RemoveRestrictedTodos background migration
+merge_request: 21893
+author:
+type: fixed
diff --git a/db/migrate/20181002172433_remove_restricted_todos_with_cte.rb b/db/migrate/20181002172433_remove_restricted_todos_with_cte.rb
new file mode 100644
index 00000000000..0a8f4a12266
--- /dev/null
+++ b/db/migrate/20181002172433_remove_restricted_todos_with_cte.rb
@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+# See http://doc.gitlab.com/ce/development/migration_style_guide.html
+# for more information on how to write migrations for GitLab.
+
+# rescheduling of the revised RemoveRestrictedTodos background migration
+class RemoveRestrictedTodosWithCte < ActiveRecord::Migration
+  DOWNTIME = false
+  disable_ddl_transaction!
+
+  MIGRATION = 'RemoveRestrictedTodos'.freeze
+  BATCH_SIZE = 1000
+  DELAY_INTERVAL = 5.minutes.to_i
+
+  class Project < ActiveRecord::Base
+    include EachBatch
+
+    self.table_name = 'projects'
+  end
+
+  def up
+    Project.where('EXISTS (SELECT 1 FROM todos WHERE todos.project_id = projects.id)')
+      .each_batch(of: BATCH_SIZE) do |batch, index|
+      range = batch.pluck('MIN(id)', 'MAX(id)').first
+
+      BackgroundMigrationWorker.perform_in(index * DELAY_INTERVAL, MIGRATION, range)
+    end
+  end
+
+  def down
+    # nothing to do
+  end
+end
diff --git a/db/schema.rb b/db/schema.rb
index 773bfb96b93..4ff0272428a 100644
--- a/db/schema.rb
+++ b/db/schema.rb
@@ -11,7 +11,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 20180924201039) do
+ActiveRecord::Schema.define(version: 20181002172433) do
 
   # These are extensions that must be enabled in order to support this database
   enable_extension "plpgsql"
diff --git a/lib/gitlab/background_migration/remove_restricted_todos.rb b/lib/gitlab/background_migration/remove_restricted_todos.rb
index 68f3fa62170..9941c2fe6d9 100644
--- a/lib/gitlab/background_migration/remove_restricted_todos.rb
+++ b/lib/gitlab/background_migration/remove_restricted_todos.rb
@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 # rubocop:disable Style/Documentation
+# rubocop:disable Metrics/ClassLength
 
 module Gitlab
   module BackgroundMigration
@@ -49,11 +50,14 @@ module Gitlab
       private
 
       def remove_non_members_todos(project_id)
-        Todo.where(project_id: project_id)
-          .where('user_id NOT IN (?)', authorized_users(project_id))
-          .each_batch(of: 5000) do |batch|
-            batch.delete_all
-          end
+        if Gitlab::Database.postgresql?
+          batch_remove_todos_cte(project_id)
+        else
+          unauthorized_project_todos(project_id)
+            .each_batch(of: 5000) do |batch|
+              batch.delete_all
+            end
+        end
       end
 
       def remove_confidential_issue_todos(project_id)
@@ -86,10 +90,13 @@ module Gitlab
 
           next if target_types.empty?
 
-          Todo.where(project_id: project_id)
-            .where('user_id NOT IN (?)', authorized_users(project_id))
-            .where(target_type: target_types)
-            .delete_all
+          if Gitlab::Database.postgresql?
+            batch_remove_todos_cte(project_id, target_types)
+          else
+            unauthorized_project_todos(project_id)
+              .where(target_type: target_types)
+              .delete_all
+          end
         end
       end
 
@@ -100,6 +107,65 @@ module Gitlab
       def authorized_users(project_id)
         ProjectAuthorization.select(:user_id).where(project_id: project_id)
       end
+
+      def unauthorized_project_todos(project_id)
+        Todo.where(project_id: project_id)
+          .where('user_id NOT IN (?)', authorized_users(project_id))
+      end
+
+      def batch_remove_todos_cte(project_id, target_types = nil)
+        loop do
+          count = remove_todos_cte(project_id, target_types)
+
+          break if count == 0
+        end
+      end
+
+      def remove_todos_cte(project_id, target_types = nil)
+        sql = []
+        sql << with_all_todos_sql(project_id, target_types)
+        sql << as_deleted_sql
+        sql << "SELECT count(*) FROM deleted"
+
+        result = Todo.connection.exec_query(sql.join(' '))
+        result.rows[0][0].to_i
+      end
+
+      def with_all_todos_sql(project_id, target_types = nil)
+        if target_types
+          table = Arel::Table.new(:todos)
+          in_target = table[:target_type].in(target_types)
+          target_types_sql = " AND #{in_target.to_sql}"
+        end
+
+        <<-SQL
+          WITH all_todos AS (
+          	SELECT id
+          	FROM "todos"
+          	WHERE "todos"."project_id" = #{project_id}
+            AND (user_id NOT IN (
+              SELECT "project_authorizations"."user_id"
+              FROM "project_authorizations"
+              WHERE "project_authorizations"."project_id" = #{project_id})
+              #{target_types_sql}
+            )
+        	),
+        SQL
+      end
+
+      def as_deleted_sql
+        <<-SQL
+          deleted AS (
+            DELETE FROM todos
+            WHERE id IN (
+              SELECT id
+              FROM all_todos
+              LIMIT 5000
+            )
+            RETURNING id
+          )
+        SQL
+      end
     end
   end
 end