Defer tick_total calculation to the background data migration starting time instead of enqueueing time

Author: fatkodima

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## master (unreleased)
 
 - Drop support for Ruby < 3.3
+- Defer `tick_total` calculation to the background data migration starting time instead of enqueueing time
 - Make "replacing an index" check independent of the live database state
 
 ## 0.34.0 (2026-05-29)

lib/online_migrations/background_data_migrations/migration.rb

@@ -123,7 +123,13 @@ def stuck?
       # @private
       def start
         if enqueued?
-          update!(status: :running, started_at: Time.current)
+          # Defer `tick_total` calculation to the migration starting time
+          # instead of the enqueueing time to avoid failed deploys.
+          self.tick_total ||= safely_calculate_tick_total
+          self.status = :running
+          self.started_at = Time.current
+          save!
+
           data_migration.after_start
           true
         else
@@ -309,11 +315,18 @@ def retry
         def set_defaults
           config = ::OnlineMigrations.config.background_data_migrations
           self.max_attempts ||= config.max_attempts
-          self.tick_total ||= on_shard_if_present do
+          self.iteration_pause ||= config.iteration_pause
+        end
+
+        def safely_calculate_tick_total
+          on_shard_if_present do
             data_migration.count
           end
-
-          self.iteration_pause ||= config.iteration_pause
+        rescue ActiveRecord::QueryCanceled
+          # `tick_total` is not required and is used only for progress tracking.
+          # Probably the `count` method was implemented in a non-efficient way.
+          # Better to not track progress than have a failing migration.
+          nil
         end
 
         def instrument_status_change

test/background_data_migrations/data_migrations.rb

@@ -192,4 +192,10 @@ def count
       collection.count
     end
   end
+
+  class HeavyCountMigration < WithCountMigration
+    def count
+      ActiveRecord::Base.connection.execute("select pg_sleep(100)")
+    end
+  end
 end

test/background_data_migrations/migration_job_test.rb

@@ -106,6 +106,16 @@ def test_uses_count_method_if_present
       assert_equal 3, m.reload.tick_total
     end
 
+    def test_gracefully_handles_heavy_count_method
+      @connection.execute("set statement_timeout to '0.01s'")
+      m = create_migration("HeavyCountMigration")
+      MigrationJob.perform_inline(m.id)
+
+      assert_nil m.tick_total
+    ensure
+      @connection.execute("set statement_timeout to '10s'")
+    end
+
     def test_stores_metadada_about_the_data_migration
       m = create_migration("ArrayCollectionMigration")
       MigrationJob.perform_inline(m.id)

test/background_data_migrations/migration_test.rb

@@ -37,6 +37,7 @@ def test_sets_defaults
         config.stub(:iteration_pause, 10) do
           m = create_migration
           assert m.pending?
+          assert_nil m.tick_total
           assert_equal config.max_attempts, m.max_attempts
           assert_equal 10, m.iteration_pause
         end
@@ -55,7 +56,8 @@ def test_progress_succeeded_migration
 
     def test_progress_not_finished_migration
       2.times { User.create! }
-      m = create_migration(migration_name: "MigrationWithCount")
+      # tick_total is calculated on migration run time, so need to set explicitly.
+      m = create_migration(migration_name: "MigrationWithCount", tick_total: 2)
 
       m.update!(tick_count: 1)
       assert_in_delta 50.0, m.progress
@@ -66,7 +68,8 @@ def test_progress_not_finished_migration
 
     def test_progress_running_migration_without_records
       assert_equal 0, EmptyCollection.new.count
-      m = create_migration(migration_name: "EmptyCollection")
+      # tick_total is calculated on migration run time, so need to set explicitly.
+      m = create_migration(migration_name: "EmptyCollection", tick_total: 0)
       assert_in_delta 0.0, m.progress
     end