Internal improvements 2

db/migrate/20200702202022_create_active_storage_db_files.rb

@@ -20,7 +20,6 @@ def change
 
   def primary_key_type
     config = Rails.configuration.generators
-    primary_key_type = config.options[config.orm][:primary_key_type]
-    primary_key_type || :primary_key
+    config.options.dig(config.orm, :primary_key_type) || :primary_key
   end
 end

lib/active_storage/service/db_service.rb

@@ -18,8 +18,10 @@ class Service::DBService < Service
     end
     # :nocov:
 
+    MINIMUM_CHUNK_SIZE = 1
+
     def initialize(public: false, **)
-      @chunk_size = ENV.fetch("ASDB_CHUNK_SIZE") { 1.megabytes }.to_i
+      @chunk_size = [ENV.fetch("ASDB_CHUNK_SIZE") { 1.megabytes }.to_i, MINIMUM_CHUNK_SIZE].max
       @public = public
     end
 
@@ -49,7 +51,7 @@ def download_chunk(key, range)
         size = range.size
         args = adapter_sqlserver? || adapter_sqlite? ? "data, #{from}, #{size}" : "data FROM #{from} FOR #{size}"
         record = object_for(key, fields: "SUBSTRING(#{args}) AS chunk")
-        raise(ActiveStorage::FileNotFoundError) unless record
+        raise ActiveStorage::FileNotFoundError unless record
 
         record.chunk
       end
@@ -158,15 +160,18 @@ def generate_url(key, expires_in:, filename:, content_type:, disposition:)
     end
 
     def ensure_integrity_of(key, checksum)
-      return if Digest::MD5.base64digest(object_for(key).data) == 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
+      raise ActiveStorage::FileNotFoundError unless file
 
       file.data
     end

spec/dummy/app/controllers/posts_controller.rb

@@ -72,6 +72,6 @@ def post_params
   end
 
   def unprocessable
-    Gem::Version.new(Rails.version) >= Gem::Version.new("7.1") ? :unprocessable_content : :unprocessable_entity
+    Rails::VERSION::MAJOR > 7 || (Rails::VERSION::MAJOR == 7 && Rails::VERSION::MINOR >= 1) ? :unprocessable_content : :unprocessable_entity
   end
 end

spec/requests/file_controller_spec.rb

@@ -2,7 +2,7 @@
 
 RSpec.describe "File controller" do
   def unprocessable
-    Gem::Version.new(Rails.version) >= Gem::Version.new("7.1") ? :unprocessable_content : :unprocessable_entity
+    Rails::VERSION::MAJOR > 7 || (Rails::VERSION::MAJOR == 7 && Rails::VERSION::MINOR >= 1) ? :unprocessable_content : :unprocessable_entity
   end
 
   let(:blob) { create_blob(filename: "img.jpg", content_type: "image/jpeg") }

spec/service/active_storage/service/db_service_spec.rb

@@ -77,6 +77,31 @@
         expect(compose).to be_a ActiveStorageDB::File
         expect(compose.data).to eq [first_db_file.data, second_db_file.data, third_db_file.data].join
       end
+
+      context "when a source key is missing" do
+        subject(:compose) { service.compose(%w[key1 missing_key key3], "dest_key") }
+
+        it "raises FileNotFoundError" do
+          expect { compose }.to raise_exception(ActiveStorage::FileNotFoundError)
+        end
+
+        it "does not create the destination file" do
+          compose rescue nil # rubocop:disable Style/RescueModifier
+          expect(ActiveStorageDB::File.find_by(ref: "dest_key")).to be_nil
+        end
+      end
+
+      context "with empty source keys" do
+        subject(:compose) { service.compose([], "dest_key") }
+
+        it "does not create a destination file" do
+          expect { compose }.not_to change(ActiveStorageDB::File, :count)
+        end
+
+        it "returns nil" do
+          expect(compose).to be_nil
+        end
+      end
     end
   end
 
@@ -238,7 +263,7 @@
   end
 
   describe ".url_for_direct_upload" do
-    subject do
+    subject(:url_for_direct_upload) do
       service.url_for_direct_upload(
         key,
         expires_in: 5.minutes,
@@ -260,5 +285,49 @@
     after { service.delete(key) }
 
     it { is_expected.to start_with "#{url_options[:protocol]}#{url_options[:host]}" }
+
+    it "generates a token that contains the expected fields" do
+      url = url_for_direct_upload
+      token = url.split("/").last
+      verified = ActiveStorage.verifier.verified(token, purpose: :blob_token).symbolize_keys
+      expect(verified).to include(
+        key: key,
+        content_type: content_type,
+        content_length: fixture_data.size,
+        checksum: checksum
+      )
+      expect(verified[:service_name]).to be_present
+    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 "chunk_size validation" do
+    it "enforces a minimum chunk size of 1" do
+      allow(ENV).to receive(:fetch).with("ASDB_CHUNK_SIZE").and_return("0")
+      svc = described_class.configure(:tmp, tmp: { service: "DB" })
+      expect(svc.instance_variable_get(:@chunk_size)).to eq(1)
+    end
+
+    it "enforces a minimum chunk size for negative values" do
+      allow(ENV).to receive(:fetch).with("ASDB_CHUNK_SIZE").and_return("-5")
+      svc = described_class.configure(:tmp, tmp: { service: "DB" })
+      expect(svc.instance_variable_get(:@chunk_size)).to eq(1)
+    end
   end
 end

spec/support/shared_contexts/rake_context.rb

@@ -1,7 +1,11 @@
 # frozen_string_literal: true
 
-shared_context 'with rake tasks' do
-  Rails.application.load_tasks
+require "rake"
+
+shared_context "with rake tasks" do
+  before(:context) do
+    Rails.application.load_tasks unless Rake::Task.task_defined?("environment")
+  end
 
   def execute_task(task, args = nil)
     with_captured_stdout do