Adjust some internal methods

Author: Mattia Roccoberton

lib/active_storage/service/db_service.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 module ActiveStorage
+  # Wraps a DB table as an Active Storage service. See ActiveStorage::Service
+  # for the generic API documentation that applies to all services.
   class Service::DBService < Service
     def initialize(**)
       @chunk_size = ENV.fetch('ASDB_CHUNK_SIZE') { 1.megabytes }
@@ -22,22 +24,27 @@ def download(key, &block)
       else
         instrument :download, key: key do
           record = ::ActiveStorageDB::File.find_by(ref: key)
-          record&.data || raise(ActiveStorage::FileNotFoundError)
+          raise(ActiveStorage::FileNotFoundError) unless record
+
+          record.data
         end
       end
     end
 
     def download_chunk(key, range)
       instrument :download_chunk, key: key, range: range do
         chunk_select = "SUBSTRING(data FROM #{range.begin + 1} FOR #{range.size}) AS chunk"
-        ::ActiveStorageDB::File.select(chunk_select).find_by(ref: key)&.chunk ||
-          raise(ActiveStorage::FileNotFoundError)
+        record = ::ActiveStorageDB::File.select(chunk_select).find_by(ref: key)
+        raise(ActiveStorage::FileNotFoundError) unless record
+
+        record.chunk
       end
     end
 
     def delete(key)
       instrument :delete, key: key do
         ::ActiveStorageDB::File.find_by(ref: key)&.destroy
+        # Ignore files already deleted
       end
     end
 
@@ -55,31 +62,14 @@ def exist?(key)
       end
     end
 
-    def url(key, expires_in:, filename:, disposition:, content_type:)
-      instrument :url, key: key do |payload|
-        content_disposition = content_disposition_with(type: disposition, filename: filename)
-        verified_key_with_expiration = ActiveStorage.verifier.generate(
-          {
-            key: key,
-            disposition: content_disposition,
-            content_type: content_type
-          },
-          expires_in: expires_in,
-          purpose: :blob_key
-        )
-        current_uri = URI.parse(current_host)
-        generated_url = url_helpers.service_url(
-          verified_key_with_expiration,
-          protocol: current_uri.scheme,
-          host: current_uri.host,
-          port: current_uri.port,
-          disposition: content_disposition,
-          content_type: content_type,
-          filename: filename
-        )
-        payload[:url] = generated_url
-
-        generated_url
+    if Rails::VERSION::MAJOR == 6 && Rails::VERSION::MINOR == 0
+      def url(key, expires_in:, filename:, disposition:, content_type:)
+        instrument :url, key: key do |payload|
+          opts = { expires_in: expires_in, filename: filename, content_type: content_type, disposition: disposition }
+          generate_url(key, opts).tap do |generated_url|
+            payload[:url] = generated_url
+          end
+        end
       end
     end
 
@@ -90,15 +80,16 @@ def url_for_direct_upload(key, expires_in:, content_type:, content_length:, chec
             key: key,
             content_type: content_type,
             content_length: content_length,
-            checksum: checksum
+            checksum: checksum,
+            service_name: respond_to?(:name) ? name : 'db'
           },
           expires_in: expires_in,
           purpose: :blob_token
         )
-        generated_url = url_helpers.update_service_url(verified_token_with_expiration, host: current_host)
-        payload[:url] = generated_url
 
-        generated_url
+        url_helpers.update_service_url(verified_token_with_expiration, url_options).tap do |generated_url|
+          payload[:url] = generated_url
+        end
       end
     end
 
@@ -119,6 +110,39 @@ def current_host
       end
     end
 
+    def private_url(key, expires_in:, filename:, content_type:, disposition:, **)
+      generate_url(key, expires_in: expires_in, filename: filename, content_type: content_type, disposition: disposition)
+    end
+
+    def public_url(key, filename:, content_type: nil, disposition: :attachment, **)
+      generate_url(key, expires_in: nil, filename: filename, content_type: content_type, disposition: disposition)
+    end
+
+    def generate_url(key, expires_in:, filename:, content_type:, disposition:)
+      content_disposition = content_disposition_with(type: disposition, filename: filename)
+      verified_key_with_expiration = ActiveStorage.verifier.generate(
+        {
+          key: key,
+          disposition: content_disposition,
+          content_type: content_type,
+          service_name: respond_to?(:name) ? name : 'db'
+        },
+        expires_in: expires_in,
+        purpose: :blob_key
+      )
+
+      current_uri = URI.parse(current_host)
+      url_helpers.service_url(
+        verified_key_with_expiration,
+        protocol: current_uri.scheme,
+        host: current_uri.host,
+        port: current_uri.port,
+        disposition: content_disposition,
+        content_type: content_type,
+        filename: filename
+      )
+    end
+
     def ensure_integrity_of(key, checksum)
       file = ::ActiveStorageDB::File.find_by(ref: key)
       return if Digest::MD5.base64digest(file.data) == checksum
@@ -140,5 +164,13 @@ def stream(key)
     def url_helpers
       @url_helpers ||= ::ActiveStorageDB::Engine.routes.url_helpers
     end
+
+    def url_options
+      if ActiveStorage::Current.respond_to? :url_options
+        ActiveStorage::Current.url_options
+      else
+        { host: ActiveStorage::Current.host }
+      end
+    end
   end
 end

spec/requests/file_controller_spec.rb

@@ -21,14 +21,14 @@ def create_blob_before_direct_upload(byte_size:, checksum:, filename: 'hello.txt
   end
 
   let(:blob) { create_blob(filename: 'img.jpg', content_type: 'image/jpg') }
-  let(:host) { 'http://test.example.com:3001' }
   let(:url_options) do
     {
       protocol: 'http://',
       host: 'test.example.com',
       port: '3001',
     }
   end
+  let(:host) { "#{url_options[:protocol]}#{url_options[:host]}:#{url_options[:port]}" }
   let(:engine_url_helpers) { ::ActiveStorageDB::Engine.routes.url_helpers }
 
   before do
@@ -137,5 +137,19 @@ def create_blob_before_direct_upload(byte_size:, checksum:, filename: 'hello.txt
         expect(response).to have_http_status(:not_found)
       end
     end
+
+    context 'when the integrity check fails' do
+      let(:invalid_file) { create(:active_storage_db_file, data: 'Some other data') }
+
+      before do
+        allow(::ActiveStorageDB::File).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_entity)
+      end
+    end
   end
 end

spec/service/active_storage/service/db_service_spec.rb

@@ -3,14 +3,14 @@
 RSpec.describe ActiveStorage::Service::DBService do
   let(:fixture_data) { build(:active_storage_db_file).data }
   let(:content_type) { 'image/png' }
-  let(:host) { 'http://test.example.com:3001' }
   let(:url_options) do
     {
       protocol: 'http://',
       host: 'test.example.com',
       port: '3001',
     }
   end
+  let(:host) { "#{url_options[:protocol]}#{url_options[:host]}:#{url_options[:port]}" }
 
   let(:checksum) { Digest::MD5.base64digest(fixture_data) }
   let(:key) { SecureRandom.base58(24) }
@@ -22,6 +22,32 @@
     allow(ENV).to receive(:fetch).with('ASDB_CHUNK_SIZE').and_return(100)
   end
 
+  describe 'public URL generation' do
+    let(:config) do
+      {
+        db: {
+          service: 'db',
+          public: true
+        }
+      }
+    end
+
+    before do
+      if ActiveStorage::Current.respond_to? :url_options
+        allow(ActiveStorage::Current).to receive(:url_options).and_return(url_options)
+      else
+        allow(ActiveStorage::Current).to receive(:host).and_return(host)
+      end
+    end
+
+    it 'returns a public URL' do
+      filename = ActiveStorage::Filename.new('some_filename')
+      service = ActiveStorage::Service.configure(:db, config)
+      url = service.url('some_key', filename: filename, disposition: :inline, content_type: 'text/plain', expires_in: nil)
+      expect(url).to match %r{#{host}/active_storage_db/files/}
+    end
+  end
+
   describe '.delete' do
     subject(:delete) { service.delete(key) }
 
@@ -137,27 +163,6 @@
     end
   end
 
-  describe '.url' do
-    subject do
-      service.url(key, expires_in: 5.minutes, disposition: :inline, filename: filename, content_type: content_type)
-    end
-
-    let(:filename) { ActiveStorage::Filename.new('avatar.png') }
-
-    before do
-      upload
-      if ActiveStorage::Current.respond_to? :url_options
-        allow(ActiveStorage::Current).to receive(:url_options).and_return(url_options)
-      else
-        allow(ActiveStorage::Current).to receive(:host).and_return(host)
-      end
-    end
-
-    after { service.delete(key) }
-
-    it { is_expected.to start_with "#{url_options[:protocol]}#{url_options[:host]}" }
-  end
-
   describe '.url_for_direct_upload' do
     subject do
       service.url_for_direct_upload(