feat: Improve upload checksum

Compute checksum from in-memory IO data before persisting, instead of writing first and re-reading the entire blob from the DB. Remove the now-dead ensure_integrity_of method.

Author: blocknotes

lib/active_storage/service/db_service.rb

@@ -27,9 +27,12 @@ def initialize(public: false, **)
 
     def upload(key, io, checksum: nil, **)
       instrument :upload, key: key, checksum: checksum do
-        file = ::ActiveStorageDB::File.create!(ref: key, data: io.read)
-        ensure_integrity_of(key, checksum) if checksum
-        file
+        data = io.read
+        if checksum
+          digest = Digest::MD5.base64digest(data)
+          raise ActiveStorage::IntegrityError unless digest == checksum
+        end
+        ::ActiveStorageDB::File.create!(ref: key, data: data)
       end
     end
 
@@ -159,16 +162,6 @@ def generate_url(key, expires_in:, filename:, content_type:, disposition:)
       )
     end
 
-    def ensure_integrity_of(key, checksum)
-      record = object_for(key)
-      raise ActiveStorage::FileNotFoundError unless record
-
-      return if Digest::MD5.base64digest(record.data) == checksum
-
-      delete(key)
-      raise ActiveStorage::IntegrityError
-    end
-
     def retrieve_file(key)
       file = object_for(key)
       raise ActiveStorage::FileNotFoundError unless file

spec/requests/file_controller_spec.rb

@@ -143,18 +143,13 @@ def unprocessable
     end
 
     context "when the integrity check fails" do
-      let(:invalid_file) { create(:active_storage_db_file, data: "Some other data") }
+      it "fails to upload and does not persist the file", :aggregate_failures do
+        allow(blob.service).to receive(:upload).and_raise(ActiveStorage::IntegrityError)
 
-      before do
-        annotated_scope = ActiveStorageDB::File.annotate("DBService#object_for")
-        allow(ActiveStorageDB::File).to receive(:annotate).and_return(annotated_scope)
-        allow(annotated_scope).to receive(:find_by).and_return(invalid_file)
-      end
-
-      it "fails to upload" do
         put blob.service_url_for_direct_upload, params: data, headers: { "Content-Type" => "text/plain" }
 
         expect(response).to have_http_status(unprocessable)
+        expect(blob.service).not_to exist(blob.key)
       end
     end
   end

spec/service/active_storage/service/db_service_spec.rb

@@ -247,6 +247,10 @@
       it "fails to upload the data" do
         expect { upload }.to raise_exception ActiveStorage::IntegrityError
       end
+
+      it "does not persist the file" do
+        expect { upload rescue nil }.not_to change(ActiveStorageDB::File, :count) # rubocop:disable Style/RescueModifier
+      end
     end
 
     context "with a duplicate key" do
@@ -300,23 +304,6 @@
     end
   end
 
-  describe ".ensure_integrity_of" do
-    before { upload }
-
-    after { service.delete(key) rescue nil } # rubocop:disable Style/RescueModifier
-
-    context "when the record is missing during integrity check" do
-      it "raises FileNotFoundError" do
-        # Simulate the record disappearing between upload and integrity check
-        allow(service).to receive(:object_for).with(key).and_return(nil)
-
-        expect {
-          service.send(:ensure_integrity_of, key, "bad_checksum")
-        }.to raise_exception(ActiveStorage::FileNotFoundError)
-      end
-    end
-  end
-
   describe "service_name_for_token" do
     it "returns the service name when set" do
       service.name = "custom_db"