From 4f761032c7688c62d168cb9d1edf92eea8d5a341 Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Thu, 31 Oct 2024 23:20:03 +0200
Subject: [PATCH] Fix background data migrations to work with
 `includes`/`eager_load` on relation

Fixes #65.
---
 CHANGELOG.md                                  |  1 +
 lib/online_migrations/batch_iterator.rb       |  3 ++-
 .../background_migrations.rb                  | 10 ++++++++++
 test/background_migrations/migration_test.rb  | 19 ++++++++++++++++---
 test/support/models.rb                        |  7 +++++++
 5 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ae52daf..febb34c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## master (unreleased)
 
+- Fix background data migrations to work with `includes`/`eager_load` on relation
 - Fix problem with running migrations for `activerecord` 8.0.0.rc2
 
 ## 0.20.0 (2024-10-21)
diff --git a/lib/online_migrations/batch_iterator.rb b/lib/online_migrations/batch_iterator.rb
index 555876c..242782c 100644
--- a/lib/online_migrations/batch_iterator.rb
+++ b/lib/online_migrations/batch_iterator.rb
@@ -19,7 +19,8 @@ def each_batch(of: 1000, column: relation.primary_key, start: nil, finish: nil,
       end
 
       relation = apply_limits(self.relation, column, start, finish, order)
-      base_relation = relation.reselect(column).reorder(column => order)
+      unscopes = Utils.ar_version < 7.1 ? [:includes] : [:includes, :preload, :eager_load]
+      base_relation = relation.unscope(*unscopes).reselect(column).reorder(column => order)
 
       start_id = start || begin
         start_row = base_relation.uncached { base_relation.first }
diff --git a/test/background_migrations/background_migrations.rb b/test/background_migrations/background_migrations.rb
index 0607aa3..caa46fa 100644
--- a/test/background_migrations/background_migrations.rb
+++ b/test/background_migrations/background_migrations.rb
@@ -91,6 +91,16 @@ def count
     end
   end
 
+  class RelationWithIncludes < OnlineMigrations::BackgroundMigration
+    def relation
+      Project.includes(user: :commits)
+    end
+
+    def process_batch(_users)
+      # no-op
+    end
+  end
+
   class MigrationWithCount < OnlineMigrations::BackgroundMigration
     def relation
       User.all
diff --git a/test/background_migrations/migration_test.rb b/test/background_migrations/migration_test.rb
index c6945d2..a00e90a 100644
--- a/test/background_migrations/migration_test.rb
+++ b/test/background_migrations/migration_test.rb
@@ -4,9 +4,6 @@
 
 module BackgroundMigrations
   class MigrationTest < Minitest::Test
-    class User < ActiveRecord::Base
-    end
-
     def setup
       @connection = ActiveRecord::Base.connection
       @connection.create_table(:users, force: true) do |t|
@@ -18,10 +15,17 @@ def setup
         t.bigint :user_id
       end
 
+      @connection.create_table(:commits, force: true) do |t|
+        t.bigint :user_id
+      end
+
       User.reset_column_information
+      Project.reset_column_information
+      Commit.reset_column_information
     end
 
     def teardown
+      @connection.drop_table(:commits, if_exists: true)
       @connection.drop_table(:projects, if_exists: true)
       @connection.drop_table(:users, if_exists: true)
       OnlineMigrations::BackgroundMigrations::Migration.delete_all
@@ -155,6 +159,15 @@ def test_empty_relation
       assert m.enqueued?
     end
 
+    def test_relation_with_includes
+      user = User.create!
+      _project = user.projects.create!
+      m = create_migration(migration_name: "RelationWithIncludes")
+
+      run_all_migration_jobs(m)
+      assert m.succeeded?
+    end
+
     def test_progress_succeded_migration
       m = build_migration(status: :succeeded)
       assert_in_delta 100.0, m.progress
diff --git a/test/support/models.rb b/test/support/models.rb
index d788de1..3a18311 100644
--- a/test/support/models.rb
+++ b/test/support/models.rb
@@ -2,9 +2,16 @@
 
 class User < ActiveRecord::Base
   has_many :projects
+  has_many :commits
 end
 
 class Project < ActiveRecord::Base
+  belongs_to :user
+  has_many :commits
+end
+
+class Commit < ActiveRecord::Base
+  belongs_to :user
 end
 
 class ShardRecord < ActiveRecord::Base