From 45c364b73aa0b0c73f80d23559ca8464e68b27b0 Mon Sep 17 00:00:00 2001
From: andrea longhi <andrea@spaghetticode.it>
Date: Fri, 18 Jun 2021 15:12:08 +0200
Subject: [PATCH] Update migration tasks with batch limit

When running the tasks on large DBs it makes sense to limit the
number of records updated on each transaction in order to avoid
locking the records (or the whole table) for too long.

The default batch size is 100'000 and it can be overridden by
passing a different value when invoking the tasks. For example,
this will set the batch size to 1000 rows:

rake solidus:migrations:migrate_default_billing_addresses_to_address_book:up[1000]
---
 ...ult_billing_addresses_to_address_book.rake | 19 +++++++++++++------
 ..._billing_addresses_to_address_book_spec.rb |  4 ++--
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/core/lib/tasks/migrations/migrate_default_billing_addresses_to_address_book.rake b/core/lib/tasks/migrations/migrate_default_billing_addresses_to_address_book.rake
index 7218e0fe1c5..c193e8b664f 100644
--- a/core/lib/tasks/migrations/migrate_default_billing_addresses_to_address_book.rake
+++ b/core/lib/tasks/migrations/migrate_default_billing_addresses_to_address_book.rake
@@ -3,22 +3,29 @@
 namespace :solidus do
   namespace :migrations do
     namespace :migrate_default_billing_addresses_to_address_book do
-      task up: :environment do
-        print "Migrating default billing addresses to address book ... "
+      task :up, [:batch_size] => [:environment] do |_t, args|
+        batch_size = args[:batch_size] || 100_000
+        print "Migrating default billing addresses to address book in batches of #{batch_size} ... "
         if Spree::UserAddress.where(default_billing: true).any?
-          Spree.user_class.joins(:bill_address).update_all(bill_address_id: nil) # rubocop:disable Rails/SkipsModelValidations
+          Spree.user_class.joins(:bill_address).in_batches(of: batch_size).each do |batch|
+            batch.update_all(bill_address_id: nil) # rubocop:disable Rails/SkipsModelValidations
+          end
         end
         Spree::UserAddress.joins(
           <<~SQL
             JOIN spree_users ON spree_user_addresses.user_id = spree_users.id
                              AND spree_user_addresses.address_id = spree_users.bill_address_id
           SQL
-        ).update_all(default_billing: true)
+        ).in_batches(of: batch_size).each do |batch|
+          batch.update_all(default_billing: true) # rubocop:disable Rails/SkipsModelValidations
+        end
+
         puts "Success"
       end
 
-      task down: :environment do
-        Spree::UserAddress.update_all(default_billing: false) # rubocop:disable Rails/SkipsModelValidations
+      task :down, [:batch_size] => [:environment] do |_t, args|
+        batch_size = args[:batch_size] || 100_000
+        Spree::UserAddress.in_batches(of: batch_size).update_all(default_billing: false) # rubocop:disable Rails/SkipsModelValidations
         puts "Rolled back default billing address migration to address book"
       end
     end
diff --git a/core/spec/lib/tasks/migrations/migrate_default_billing_addresses_to_address_book_spec.rb b/core/spec/lib/tasks/migrations/migrate_default_billing_addresses_to_address_book_spec.rb
index 17982f06682..a165fd26c78 100644
--- a/core/spec/lib/tasks/migrations/migrate_default_billing_addresses_to_address_book_spec.rb
+++ b/core/spec/lib/tasks/migrations/migrate_default_billing_addresses_to_address_book_spec.rb
@@ -33,7 +33,7 @@
 
       it 'runs' do
         expect { task.invoke }.to output(
-          "Migrating default billing addresses to address book ... Success\n"
+          "Migrating default billing addresses to address book in batches of 100000 ... Success\n"
         ).to_stdout
       end
 
@@ -63,7 +63,7 @@
 
       it 'runs' do
         expect { task.invoke }.to output(
-          "Migrating default billing addresses to address book ... Success\n"
+          "Migrating default billing addresses to address book in batches of 100000 ... Success\n"
         ).to_stdout
       end