feat: WIP

Author: blocknotes

app/models/active_storage_db/file.rb

@@ -2,9 +2,7 @@
 
 module ActiveStorageDB
   class File < ApplicationRecord
-    validates :ref,
-              presence: true,
-              uniqueness: { case_sensitive: false }
+    validates :ref, presence: true, uniqueness: true
     validates :data, presence: true
   end
 end

db/migrate/20200702202022_create_active_storage_db_files.rb

@@ -13,7 +13,12 @@ def change
       end
 
       t.index [:ref], unique: true
+      t.index [:created_at]
     end
+
+    add_index :active_storage_db_files, [:ref, :created_at],
+              name: "index_active_storage_db_files_ref_created",
+              order: { created_at: :desc }
   end
 
   private

lib/active_storage/service/db_service.rb

@@ -5,23 +5,23 @@
 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
-    # :nocov:
-    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
-    # :nocov:
+    include ActiveStorage::DBServiceRails70
+
+    DEFAULT_RETRY_OPTIONS = {
+      max_attempts: 3,
+      base_delay: 0.1,
+      max_delay: 2.0,
+      retryable_errors: [
+        ActiveRecord::ConnectionFailed,
+        ActiveRecord::StatementTimeout
+      ].freeze
+    }.freeze
 
     MINIMUM_CHUNK_SIZE = 1
 
-    def initialize(public: false, **)
-      @chunk_size = [ENV.fetch("ASDB_CHUNK_SIZE") { 1.megabytes }.to_i, MINIMUM_CHUNK_SIZE].max
+    def initialize(public: false, chunk_size: nil, **)
+      @chunk_size = [chunk_size || ENV.fetch("ASDB_CHUNK_SIZE") { 1.megabyte }.to_i, MINIMUM_CHUNK_SIZE].max
       @max_size = ENV.fetch("ASDB_MAX_FILE_SIZE", nil)&.to_i
       @public = public
     end
@@ -37,7 +37,7 @@ def upload(key, io, checksum: nil, **)
           digest = Digest::MD5.base64digest(data)
           raise ActiveStorage::IntegrityError unless digest == checksum
         end
-        ::ActiveStorageDB::File.create!(ref: key, data: data)
+        retry_on_failure { ::ActiveStorageDB::File.create!(ref: key, data: data) }
       end
     end
 
@@ -55,33 +55,33 @@ def download(key, &block)
 
     def download_chunk(key, range)
       instrument :download_chunk, key: key, range: range do
-        # NOTE: from/size are derived from Range#begin and Range#size (always integers),
-        # so string interpolation into SQL is safe here.
-        from = range.begin + 1
-        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
-
-        record.chunk
+        chunk = if adapter_postgresql? && @chunk_size >= 1.megabyte
+                  pg_read_binary(key, range)
+                else
+                  sql_chunk(key, range)
+                end
+        raise ActiveStorage::FileNotFoundError unless chunk
+
+        chunk
       end
     end
 
     def delete(key)
       instrument :delete, key: key do
         comment = "DBService#delete"
-        record = ::ActiveStorageDB::File.annotate(comment).find_by(ref: key)
-        record&.destroy
-        # NOTE: Ignore files already deleted
-        !record.nil?
+        retry_on_failure do
+          ::ActiveStorageDB::File.annotate(comment).where(ref: key).delete > 0
+        end
       end
     end
 
     def delete_prefixed(prefix)
       instrument :delete_prefixed, prefix: prefix do
         comment = "DBService#delete_prefixed"
         sanitized_prefix = "#{ActiveRecord::Base.sanitize_sql_like(prefix)}%"
-        ::ActiveStorageDB::File.annotate(comment).where("ref LIKE ?", sanitized_prefix).destroy_all
+        retry_on_failure do
+          ::ActiveStorageDB::File.annotate(comment).where("ref LIKE ?", sanitized_prefix).delete_all
+        end
       end
     end
 
@@ -120,20 +120,61 @@ def headers_for_direct_upload(_key, content_type:, **)
 
     private
 
+    def retry_options
+      @retry_options ||= {
+        max_attempts: 3,
+        base_delay: 0.1,
+        max_delay: 2.0,
+        retryable_errors: default_retryable_errors
+      }
+    end
+
+    def retry_on_failure
+      attempts = 0
+      max_attempts = retry_options[:max_attempts]
+      base_delay = retry_options[:base_delay]
+      max_delay = retry_options[:max_delay]
+      retryable_errors = retry_options[:retryable_errors]
+
+      begin
+        yield
+      rescue *retryable_errors
+        attempts += 1
+        raise if attempts >= max_attempts
+
+        delay = [base_delay * (2**attempts), max_delay].min
+        sleep(delay)
+        retry
+      end
+    end
+
+    def default_retryable_errors
+      errors = [
+        ActiveRecord::ConnectionFailed,
+        ActiveRecord::StatementTimeout
+      ]
+      errors << PG::ConnectionBad if defined?(PG::ConnectionBad)
+      errors
+    end
+
     def service_name_for_token
       name.presence || "db"
     end
 
     def adapter_sqlite?
-      return @adapter_sqlite if defined?(@adapter_sqlite)
-
-      @adapter_sqlite = active_storage_db_adapter_name == "SQLite"
+      @adapter_sqlite ||= active_storage_db_adapter_name == "SQLite"
     end
 
     def adapter_sqlserver?
-      return @adapter_sqlserver if defined?(@adapter_sqlserver)
+      @adapter_sqlserver ||= active_storage_db_adapter_name == "SQLServer"
+    end
 
-      @adapter_sqlserver = active_storage_db_adapter_name == "SQLServer"
+    def adapter_postgresql?
+      @adapter_postgresql ||= active_storage_db_adapter_name == "PostgreSQL"
+    end
+
+    def adapter_mysql?
+      @adapter_mysql ||= active_storage_db_adapter_name == "Mysql2"
     end
 
     def active_storage_db_adapter_name
@@ -191,6 +232,23 @@ def stream(key)
       end
     end
 
+    def sql_chunk(key, range)
+      from = range.begin + 1
+      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")
+      record&.chunk
+    end
+
+    def pg_read_binary(key, range)
+      from = range.begin + 1
+      size = range.size
+      comment = "DBService#pg_read_binary"
+      ::ActiveStorageDB::File.annotate(comment).where(ref: key).pick("get_byte(data, (#{from} - 1) + generate_series(0, #{size} - 1))")
+    rescue ActiveRecord::StatementInvalid
+      sql_chunk(key, range)
+    end
+
     def data_size
       if adapter_sqlserver?
         "DATALENGTH(data) AS size"

lib/active_storage/service/db_service_rails70.rb

@@ -3,6 +3,16 @@
 module ActiveStorage
   module DBServiceRails70
     def compose(source_keys, destination_key, **)
+      if source_keys.length > 10 || ENV["ASDB_COMPOSE_USE_TEMP_FILE"] == "true"
+        compose_with_temp_file(source_keys, destination_key)
+      else
+        compose_in_memory(source_keys, destination_key)
+      end
+    end
+
+    private
+
+    def compose_in_memory(source_keys, destination_key)
       buffer = nil
       comment = "DBService#compose"
       source_keys.each do |source_key|
@@ -18,35 +28,21 @@ def compose(source_keys, destination_key, **)
       ::ActiveStorageDB::File.create!(ref: destination_key, data: buffer) if buffer
     end
 
-    private
-
-    def current_host
-      opts = url_options || {}
-      opts[:port] ? "#{opts[:protocol]}#{opts[:host]}:#{opts[:port]}" : "#{opts[:protocol]}#{opts[: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 compose_with_temp_file(source_keys, destination_key)
+      Tempfile.create(["active_storage_db_compose", ".bin"], binmode: true) do |tempfile|
+        comment = "DBService#compose"
+        source_keys.each do |source_key|
+          record = ::ActiveStorageDB::File.annotate(comment).find_by(ref: source_key)
+          raise ActiveStorage::FileNotFoundError unless record
 
-    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
+          tempfile.write(record.data)
+        end
+        tempfile.rewind
 
-    def url_options
-      ActiveStorage::Current.url_options
+        retry_on_failure do
+          ::ActiveStorageDB::File.create!(ref: destination_key, data: tempfile.read)
+        end
+      end
     end
   end
 end

spec/rails_helper.rb

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
-ENV['RAILS_ENV'] = 'test'
+ENV["RAILS_ENV"] = "test"
 
-require 'simplecov'
-require 'simplecov-lcov'
+require "simplecov"
+require "simplecov-lcov"
 
 SimpleCov::Formatter::LcovFormatter.config do |c|
   c.report_with_single_file = true
@@ -26,15 +26,15 @@
   add_filter %r{^/vendor/}
 end
 
-require 'spec_helper'
+require "spec_helper"
 
 require File.expand_path("dummy/config/environment.rb", __dir__)
 
-abort('The Rails environment is running in production mode!') if Rails.env.production?
-require 'rspec/rails'
-require 'factory_bot_rails'
+abort("The Rails environment is running in production mode!") if Rails.env.production?
+require "rspec/rails"
+require "factory_bot_rails"
 
-support_files = File.expand_path('support/**/*.rb', __dir__)
+support_files = File.expand_path("support/**/*.rb", __dir__)
 Dir[support_files].sort.each { |f| require f }
 
 RSpec.configure do |config|
@@ -50,14 +50,14 @@
         ActiveRecord::Base.connection_config
       end
 
-    intro = ('-' * 80)
+    intro = ("-" * 80)
     intro << "\n"
     intro << "- Ruby:          #{RUBY_VERSION}\n"
     intro << "- Rails:         #{Rails.version}\n"
     intro << "- ActiveStorage: #{ActiveStorage.version}\n"
     intro << "- DB adapter:    #{db_config[:adapter]}\n"
     intro << "- DB name:       #{db_config[:database]}\n"
-    intro << ('-' * 80)
+    intro << ("-" * 80)
 
     RSpec.configuration.reporter.message(intro)
   end