Merge pull request #17 from blocknotes/update_service_api

Mattia Roccoberton · web-flow · commit c39b9e480517 · 2022-03-02T08:25:34.000+01:00
Update Service API
diff --git a/.rubocop.yml b/.rubocop.yml
@@ -36,7 +36,7 @@ RSpec/MultipleExpectations:
 
 RSpec/MultipleMemoizedHelpers:
   # Default is 5
-  Max: 10
+  Max: 12
 
 RSpec/NestedGroups:
   # Default is 3
diff --git a/lib/active_storage/service/db_service.rb b/lib/active_storage/service/db_service.rb
@@ -1,9 +1,24 @@
 # frozen_string_literal: true
 
+require 'active_storage/service/db_service_rails60'
+require 'active_storage/service/db_service_rails61'
+require 'active_storage/service/db_service_rails70'
+
 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(**)
+    if Rails::VERSION::MAJOR >= 7
+      include ActiveStorage::DBServiceRails70
+    elsif Rails::VERSION::MAJOR == 6 && Rails::VERSION::MINOR == 1
+      include ActiveStorage::DBServiceRails61
+    else
+      include ActiveStorage::DBServiceRails60
+    end
+
+    def initialize(public: false, **)
       @chunk_size = ENV.fetch('ASDB_CHUNK_SIZE') { 1.megabytes }
+      @public = public
     end
 
     def upload(key, io, checksum: nil, **)
@@ -22,22 +37,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,50 +75,23 @@ 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
-      end
-    end
-
     def url_for_direct_upload(key, expires_in:, content_type:, content_length:, checksum:, custom_metadata: {})
       instrument :url, key: key do |payload|
         verified_token_with_expiration = ActiveStorage.verifier.generate(
           {
             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
 
@@ -108,15 +101,29 @@ def headers_for_direct_upload(_key, content_type:, **)
 
     private
 
-    def current_host
-      if ActiveStorage::Current.respond_to? :url_options
-        opts = ActiveStorage::Current.url_options || {}
-        url = "#{opts[:protocol]}#{opts[:host]}"
-        url += ":#{opts[:port]}" if opts[:port]
-        url || ''
-      else
-        ActiveStorage::Current.host
-      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)
diff --git a/lib/active_storage/service/db_service_rails60.rb b/lib/active_storage/service/db_service_rails60.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module ActiveStorage
+  module DBServiceRails60
+    def url(key, expires_in:, filename:, disposition:, content_type:)
+      instrument :url, key: key do |payload|
+        generate_url(key, expires_in: expires_in, filename: filename, content_type: content_type, disposition: disposition).tap do |generated_url|
+          payload[:url] = generated_url
+        end
+      end
+    end
+
+    private
+
+    def current_host
+      url_options[:host]
+    end
+
+    def url_options
+      {
+        host: ActiveStorage::Current.host
+      }
+    end
+  end
+end
diff --git a/lib/active_storage/service/db_service_rails61.rb b/lib/active_storage/service/db_service_rails61.rb
@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module ActiveStorage
+  module DBServiceRails61
+    private
+
+    def current_host
+      url_options[:host]
+    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 url_options
+      {
+        host: ActiveStorage::Current.host
+      }
+    end
+  end
+end
diff --git a/lib/active_storage/service/db_service_rails70.rb b/lib/active_storage/service/db_service_rails70.rb
@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module ActiveStorage
+  module DBServiceRails70
+    def compose(source_keys, destination_key, **)
+      buffer = nil
+      source_keys.each do |source_key|
+        data = ::ActiveStorageDB::File.find_by!(ref: source_key).data
+        if buffer
+          buffer << data
+        else
+          buffer = data
+        end
+      end
+      ::ActiveStorageDB::File.create!(ref: destination_key, data: buffer) if buffer
+    end
+
+    private
+
+    def current_host
+      opts = url_options || {}
+      url = "#{opts[:protocol]}#{opts[:host]}"
+      url + ":#{opts[:port]}" if opts[:port]
+    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 url_options
+      ActiveStorage::Current.url_options
+    end
+  end
+end
diff --git a/spec/rails_helper.rb b/spec/rails_helper.rb
@@ -6,6 +6,12 @@
   add_filter %r{^/lib/active_storage_db/version.rb}
   add_filter %r{^/spec/}
   add_filter %r{^/vendor/}
+
+  case ENV['RAILS']
+  when '6.0' then add_filter /_rails61|_rails70/
+  when '6.1' then add_filter /_rails60|_rails70/
+  when '7.0' then add_filter /_rails60|_rails61/
+  end
 end
 
 require 'spec_helper'
@@ -35,13 +41,22 @@
   config.use_transactional_fixtures = true
 
   config.before(:suite) do
+    db_config =
+      if ActiveRecord::Base.respond_to? :connection_db_config
+        ActiveRecord::Base.connection_db_config.configuration_hash
+      else
+        ActiveRecord::Base.connection_config
+      end
+
     intro = ('-' * 80)
     intro << "\n"
     intro << "- Ruby:          #{RUBY_VERSION}\n"
     intro << "- Rails:         #{Rails.version}\n"
     intro << "- ActiveStorage: #{ActiveStorage.version}\n"
-    intro << "- DB_TEST:       #{ENV['DB_TEST']}\n"
+    intro << "- DB adapter:    #{db_config[:adapter]}\n"
+    intro << "- DB name:       #{db_config[:database]}\n"
     intro << ('-' * 80)
+
     RSpec.configuration.reporter.message(intro)
   end
 end
diff --git a/spec/requests/file_controller_spec.rb b/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
diff --git a/spec/service/active_storage/service/db_service_spec.rb b/spec/service/active_storage/service/db_service_spec.rb
