Merge pull request #26 from blocknotes/enhancement/minor-refactorings

Minor internal improvements

Author: Mattia Roccoberton

app/controllers/active_storage_db/files_controller.rb

@@ -8,50 +8,56 @@ def show
       if (key = decode_verified_key)
         serve_file(key[:key], content_type: key[:content_type], disposition: key[:disposition])
       else
-        head :not_found
+        head(:not_found)
       end
     rescue ActiveStorage::FileNotFoundError
-      head :not_found
+      head(:not_found)
     end
 
     def update
       if (token = decode_verified_token)
-        if acceptable_content?(token)
-          db_service.upload(token[:key], request.body, checksum: token[:checksum])
-        else
-          head :unprocessable_entity
-        end
+        file_uploaded = upload_file(token, body: request.body)
+        head(file_uploaded ? :no_content : :unprocessable_entity)
       else
-        head :not_found
+        head(:not_found)
       end
     rescue ActiveStorage::IntegrityError
-      head :unprocessable_entity
+      head(:unprocessable_entity)
     end
 
     private
 
+    def acceptable_content?(token)
+      token[:content_type] == request.content_mime_type && token[:content_length] == request.content_length
+    end
+
     def db_service
       ActiveStorage::Blob.service
     end
 
     def decode_verified_key
-      ActiveStorage.verifier.verified(params[:encoded_key], purpose: :blob_key)
+      key = ActiveStorage.verifier.verified(params[:encoded_key], purpose: :blob_key)
+      key&.deep_symbolize_keys
+    end
+
+    def decode_verified_token
+      token = ActiveStorage.verifier.verified(params[:encoded_token], purpose: :blob_token)
+      token&.deep_symbolize_keys
     end
 
     def serve_file(key, content_type:, disposition:)
       options = {
         type: content_type || DEFAULT_SEND_FILE_TYPE,
         disposition: disposition || DEFAULT_SEND_FILE_DISPOSITION
       }
-      send_data db_service.download(key), options
+      send_data(db_service.download(key), options)
     end
 
-    def decode_verified_token
-      ActiveStorage.verifier.verified(params[:encoded_token], purpose: :blob_token)
-    end
+    def upload_file(token, body:)
+      return false unless acceptable_content?(token)
 
-    def acceptable_content?(token)
-      token[:content_type] == request.content_mime_type && token[:content_length] == request.content_length
+      db_service.upload(token[:key], request.body, checksum: token[:checksum])
+      true
     end
   end
 end

db/migrate/20200702202022_create_active_storage_db_files.rb

@@ -2,12 +2,25 @@
 
 class CreateActiveStorageDBFiles < ActiveRecord::Migration[6.0]
   def change
-    create_table :active_storage_db_files do |t|
+    create_table :active_storage_db_files, id: primary_key_type do |t|
       t.string :ref, null: false
       t.binary :data, null: false
-      t.datetime :created_at, null: false
+
+      if connection.supports_datetime_with_precision?
+        t.datetime :created_at, precision: 6, null: false
+      else
+        t.datetime :created_at, null: false
+      end
 
       t.index [:ref], unique: true
     end
   end
+
+  private
+
+  def primary_key_type
+    config = Rails.configuration.generators
+    primary_key_type = config.options[config.orm][:primary_key_type]
+    primary_key_type || :primary_key
+  end
 end

lib/active_storage/service/db_service.rb

@@ -36,18 +36,14 @@ def download(key, &block)
         end
       else
         instrument :download, key: key do
-          record = ::ActiveStorageDB::File.find_by(ref: key)
-          raise(ActiveStorage::FileNotFoundError) unless record
-
-          record.data
+          retrieve_file(key)
         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"
-        record = ::ActiveStorageDB::File.select(chunk_select).find_by(ref: key)
+        record = object_for(key, fields: "SUBSTRING(data FROM #{range.begin + 1} FOR #{range.size}) AS chunk")
         raise(ActiveStorage::FileNotFoundError) unless record
 
         record.chunk
@@ -63,7 +59,7 @@ def delete(key)
 
     def delete_prefixed(prefix)
       instrument :delete_prefixed, prefix: prefix do
-        ::ActiveStorageDB::File.where('ref LIKE ?', "#{prefix}%").destroy_all
+        ::ActiveStorageDB::File.where('ref LIKE ?', "#{ApplicationRecord.sanitize_sql_like(prefix)}%").destroy_all
       end
     end
 
@@ -127,17 +123,26 @@ def generate_url(key, expires_in:, filename:, content_type:, disposition:)
     end
 
     def ensure_integrity_of(key, checksum)
-      file = ::ActiveStorageDB::File.find_by(ref: key)
-      return if Digest::MD5.base64digest(file.data) == checksum
+      return if Digest::MD5.base64digest(object_for(key).data) == checksum
 
       delete(key)
       raise ActiveStorage::IntegrityError
     end
 
+    def retrieve_file(key)
+      file = object_for(key)
+      raise(ActiveStorage::FileNotFoundError) unless file
+
+      file.data
+    end
+
+    def object_for(key, fields: nil)
+      as_file = fields ? ::ActiveStorageDB::File.select(*fields) : ::ActiveStorageDB::File
+      as_file.find_by(ref: key)
+    end
+
     def stream(key)
-      size =
-        ::ActiveStorageDB::File.select('OCTET_LENGTH(data) AS size').find_by(ref: key)&.size ||
-        raise(ActiveStorage::FileNotFoundError)
+      size = object_for(key, fields: 'OCTET_LENGTH(data) AS size')&.size || raise(ActiveStorage::FileNotFoundError)
       (size / @chunk_size.to_f).ceil.times.each do |i|
         range = (i * @chunk_size..((i + 1) * @chunk_size) - 1)
         yield download_chunk(key, range)