fix(storage): make blob migration resilient

Limit storage migration to blobs on the source service, skip missing source files per blob, and keep target upload failures loud.

Clean up the service-name update flag naming and document rake task usage.

lib/active_storage/migrator.rb

@@ -2,21 +2,18 @@ require 'yaml'
 require 'erb'
 
 class ActiveStorage::Migrator
-  def self.migrate(from_service_name, to_service_name, update_service_name: true)
+  def self.migrate(from_service_name, to_service_name, should_update_service_name: true)
     configs = load_storage_config
     # Check if services are configured correctly
     if configs[from_service_name.to_s].nil? || configs[to_service_name.to_s].nil?
       raise "Error: The services '#{from_service_name}' or '#{to_service_name}' are not configured correctly."
     end
 
-    from_service = ActiveStorage::Service.configure(from_service_name, { from_service_name.to_sym => configs[from_service_name.to_s] })
-    to_service = ActiveStorage::Service.configure(to_service_name, { to_service_name.to_sym => configs[to_service_name.to_s] })
+    from_service = configure_service(from_service_name, configs)
 
     configure_blob_service(from_service)
 
-    Rails.logger.debug { "#{ActiveStorage::Blob.count} Blobs to migrate from #{from_service_name} to #{to_service_name}" }
-
-    migrate_blobs(from_service, to_service, to_service_name, update_service_name: update_service_name)
+    migrate_blobs(from_service_name, to_service_name, should_update_service_name: should_update_service_name)
   end
 
   def self.load_storage_config
@@ -28,9 +25,17 @@ class ActiveStorage::Migrator
     ActiveStorage::Blob.service = service
   end
 
-  def self.migrate_blobs(_from_service, to_service, to_service_name, update_service_name: true)
-    # Configure the blob service for the source service
-    ActiveStorage::Blob.find_each do |blob|
+  def self.configure_service(service_name, configs = load_storage_config)
+    service_config = configs[service_name.to_s]
+    ActiveStorage::Service.configure(service_name, { service_name.to_sym => service_config })
+  end
+
+  def self.migrate_blobs(from_service_name, to_service_name, should_update_service_name: true)
+    to_service = configure_service(to_service_name)
+    blobs = ActiveStorage::Blob.where(service_name: from_service_name.to_s)
+    Rails.logger.debug { "#{blobs.count} Blobs to migrate from #{from_service_name} to #{to_service_name}" }
+
+    blobs.find_each do |blob|
       Rails.logger.debug { '.' }
 
       blob.open do |io|
@@ -38,7 +43,9 @@ class ActiveStorage::Migrator
         to_service.upload(blob.key, io, checksum: checksum)
       end
 
-      blob.update!(service_name: to_service_name.to_s) if update_service_name
+      blob.update!(service_name: to_service_name.to_s) if should_update_service_name
+    rescue ActiveStorage::FileNotFoundError => e
+      Rails.logger.warn { "Skipping missing blob #{blob.id} (#{blob.key}): #{e.message}" }
     end
     Rails.logger.debug { 'Successful migration' }
   end

lib/tasks/storage_migrations.rake

@@ -1,12 +1,13 @@
 namespace :storage do
   desc 'Migrate blobs from one storage service to another'
+  # Example: FROM=local TO=amazon UPDATE_BLOB_SERVICE_NAME=true bundle exec rake storage:migrate
   task migrate: :environment do
     from_service = ENV.fetch('FROM', nil)
     to_service = ENV.fetch('TO', nil)
-    update_service_name = ActiveModel::Type::Boolean.new.cast(ENV.fetch('UPDATE_BLOB_SERVICE_NAME', true))
+    should_update_service_name = ActiveModel::Type::Boolean.new.cast(ENV.fetch('UPDATE_BLOB_SERVICE_NAME', true))
 
     raise 'Missing FROM or TO argument. Usage: FROM=service_name TO=service_name rake storage:migrate' if from_service.nil? || to_service.nil?
 
-    ActiveStorage::Migrator.migrate(from_service.to_sym, to_service.to_sym, update_service_name: update_service_name)
+    ActiveStorage::Migrator.migrate(from_service.to_sym, to_service.to_sym, should_update_service_name: should_update_service_name)
   end
 end

spec/lib/active_storage/migrator_spec.rb

@@ -14,15 +14,14 @@ RSpec.describe ActiveStorage::Migrator do
     context 'when services are configured correctly' do
       it 'migrates blobs from one service to another' do
         expect(ActiveStorage::Service).to receive(:configure).with('local', any_args)
-        expect(ActiveStorage::Service).to receive(:configure).with('amazon', any_args)
-        expect(described_class).to receive(:migrate_blobs).with(from_service_stub, to_service_stub, 'amazon', update_service_name: true)
+        expect(described_class).to receive(:migrate_blobs).with('local', 'amazon', should_update_service_name: true)
         expect { described_class.migrate('local', 'amazon') }.not_to raise_error
       end
 
       it 'passes the service-name update flag to blob migration' do
-        expect(described_class).to receive(:migrate_blobs).with(from_service_stub, to_service_stub, 'amazon', update_service_name: false)
+        expect(described_class).to receive(:migrate_blobs).with('local', 'amazon', should_update_service_name: false)
 
-        described_class.migrate('local', 'amazon', update_service_name: false)
+        described_class.migrate('local', 'amazon', should_update_service_name: false)
       end
 
       it 'does not override the application logger' do
@@ -45,28 +44,52 @@ RSpec.describe ActiveStorage::Migrator do
   end
 
   describe '.migrate_blobs' do
-    let(:from_service_stub) { instance_double(ActiveStorage::Service) }
     let(:to_service_stub) { instance_double(ActiveStorage::Service) }
-    let(:blob) { instance_double(ActiveStorage::Blob, key: 'blob-key', checksum: 'checksum') }
+    let(:source_blobs) { instance_double(ActiveRecord::Relation, count: 1) }
+    let(:blob) { instance_double(ActiveStorage::Blob, id: 1, key: 'blob-key', checksum: 'checksum') }
     let(:io) { StringIO.new('blob-content') }
 
     before do
-      allow(ActiveStorage::Blob).to receive(:find_each).and_yield(blob)
+      allow(ActiveStorage::Service).to receive(:configure).with(:amazon, any_args).and_return(to_service_stub)
+      allow(ActiveStorage::Blob).to receive(:where).with(service_name: 'local').and_return(source_blobs)
+      allow(source_blobs).to receive(:find_each).and_yield(blob)
       allow(blob).to receive(:open).and_yield(io)
+      allow(blob).to receive(:update!)
+      allow(to_service_stub).to receive(:upload)
     end
 
     it 'migrates blobs regardless of content type' do
       expect(to_service_stub).to receive(:upload).with('blob-key', io, checksum: 'checksum').ordered
       expect(blob).to receive(:update!).with(service_name: 'amazon').ordered
 
-      described_class.migrate_blobs(from_service_stub, to_service_stub, :amazon, update_service_name: true)
+      described_class.migrate_blobs(:local, :amazon, should_update_service_name: true)
+    end
+
+    it 'only migrates blobs from the source service' do
+      expect(ActiveStorage::Blob).to receive(:where).with(service_name: 'local').and_return(source_blobs)
+      expect(source_blobs).to receive(:find_each).and_yield(blob)
+
+      described_class.migrate_blobs(:local, :amazon, should_update_service_name: true)
     end
 
     it 'skips the blob service update when the flag is disabled' do
       expect(to_service_stub).to receive(:upload).with('blob-key', io, checksum: 'checksum')
       expect(blob).not_to receive(:update!)
 
-      described_class.migrate_blobs(from_service_stub, to_service_stub, :amazon, update_service_name: false)
+      described_class.migrate_blobs(:local, :amazon, should_update_service_name: false)
+    end
+
+    it 'skips missing source files and continues migration' do
+      missing_blob = instance_double(ActiveStorage::Blob, id: 2, key: 'missing-key')
+
+      allow(source_blobs).to receive(:find_each).and_yield(missing_blob).and_yield(blob)
+      allow(missing_blob).to receive(:open).and_raise(ActiveStorage::FileNotFoundError)
+
+      expect(missing_blob).not_to receive(:update!)
+      expect(to_service_stub).to receive(:upload).with('blob-key', io, checksum: 'checksum')
+      expect(blob).to receive(:update!).with(service_name: 'amazon')
+
+      described_class.migrate_blobs(:local, :amazon, should_update_service_name: true)
     end
 
     it 'does not update the blob service when upload fails' do
@@ -74,7 +97,7 @@ RSpec.describe ActiveStorage::Migrator do
 
       expect(blob).not_to receive(:update!)
       expect do
-        described_class.migrate_blobs(from_service_stub, to_service_stub, :amazon, update_service_name: true)
+        described_class.migrate_blobs(:local, :amazon, should_update_service_name: true)
       end.to raise_error(ActiveStorage::IntegrityError)
     end
   end

spec/lib/tasks/rake/task_storage_migrations_spec.rb

@@ -34,7 +34,7 @@ RSpec.describe Rake::Task do
 
       context 'when required arguments are present' do
         it 'updates blob service names by default' do
-          expect(ActiveStorage::Migrator).to receive(:migrate).with(:local, :amazon, update_service_name: true)
+          expect(ActiveStorage::Migrator).to receive(:migrate).with(:local, :amazon, should_update_service_name: true)
 
           with_modified_env FROM: 'local', TO: 'amazon', UPDATE_BLOB_SERVICE_NAME: nil do
             task.invoke
@@ -42,7 +42,7 @@ RSpec.describe Rake::Task do
         end
 
         it 'skips updating blob service names when disabled' do
-          expect(ActiveStorage::Migrator).to receive(:migrate).with(:local, :amazon, update_service_name: false)
+          expect(ActiveStorage::Migrator).to receive(:migrate).with(:local, :amazon, should_update_service_name: false)
 
           with_modified_env FROM: 'local', TO: 'amazon', UPDATE_BLOB_SERVICE_NAME: 'false' do
             task.invoke