WIP: Use Retrievers for fetching remote / local data

.github/workflows/ci.yml

@@ -22,5 +22,5 @@ jobs:
       with:
         ruby-version: ${{ matrix.ruby-version }}
         bundler-cache: true # 'bundle install' and cache
-    - name: Run tests
+    - name: Run tests and standard
       run: bundle exec rake

.standard.yml

@@ -0,0 +1,4 @@
+ruby_version: 2.7
+ignore:
+  - 'spec/**/*':
+    - Lint/ConstantDefinitionInBlock

Gemfile

@@ -1,4 +1,4 @@
-source 'https://rubygems.org'
+source "https://rubygems.org"
 
 # Specify your gem's dependencies in zipline.gemspec
 gemspec

Rakefile

@@ -1,12 +1,8 @@
 #!/usr/bin/env rake
 require "bundler/gem_tasks"
+require "standard/rake"
+require "rspec/core/rake_task"
 
-begin
-  require 'rspec/core/rake_task'
+RSpec::Core::RakeTask.new(:spec)
 
-  RSpec::Core::RakeTask.new(:spec)
-
-  task default: :spec
-rescue LoadError
-  # no rspec available
-end
+task default: [:spec, :standard]

lib/zipline.rb

@@ -1,7 +1,4 @@
-require 'content_disposition'
-require 'zip_kit'
-require 'zipline/version'
-require 'zipline/zip_handler'
+require "zip_kit"
 
 # class MyController < ApplicationController
 #   include Zipline
@@ -12,12 +9,18 @@
 #   end
 # end
 module Zipline
+  require_relative "zipline/version"
+  require_relative "zipline/zip_handler"
+  Dir.glob(__dir__ + "/zipline/retrievers/*.rb").sort.each do |retriever_module_path|
+    require retriever_module_path
+  end
+
   def self.included(into_controller)
     into_controller.include(ZipKit::RailsStreaming)
     super
   end
 
-  def zipline(files, zipname = 'zipline.zip', **kwargs_for_zip_kit_stream)
+  def zipline(files, zipname = "zipline.zip", **kwargs_for_zip_kit_stream)
     zip_kit_stream(filename: zipname, **kwargs_for_zip_kit_stream) do |zip_kit_streamer|
       handler = Zipline::ZipHandler.new(zip_kit_streamer, logger)
       files.each do |file, name, options = {}|

lib/zipline/retrievers.rb

@@ -0,0 +1 @@
+

lib/zipline/retrievers/active_storage_retriever.rb

@@ -0,0 +1,34 @@
+class Zipline::ActiveStorageRetriever
+  def self.build_for(item)
+    return unless defined?(ActiveStorage)
+    return new(item.blob) if is_active_storage_attachment?(item) || is_active_storage_one?(item)
+    return new(item) if is_active_storage_blob?(item)
+    nil
+  end
+
+  def self.is_active_storage_attachment?(item)
+    item.is_a?(ActiveStorage::Attachment)
+  end
+
+  def self.is_active_storage_one?(item)
+    item.is_a?(ActiveStorage::Attached::One)
+  end
+
+  def self.is_active_storage_blob?(item)
+    item.is_a?(ActiveStorage::Blob)
+  end
+
+  def initialize(blob)
+    @blob = blob
+  end
+
+  def retrieve_into(destination)
+    @blob.download do |bytes|
+      destination.write(bytes)
+    end
+  end
+
+  def restartable?(_exception)
+    false
+  end
+end

lib/zipline/retrievers/carrierwave_retriever.rb

@@ -0,0 +1,7 @@
+class Zipline::CarrierwaveRetriever
+  def self.build_for(item)
+    if defined?(CarrierWave::Storage::Fog::File) && item.is_a?(CarrierWave::Storage::Fog::File)
+      Zipline::HTTPRetriever.new(item.url)
+    end
+  end
+end

lib/zipline/retrievers/file_retriever.rb

@@ -0,0 +1,21 @@
+class Zipline::FileRetriever
+  def self.build_for(item)
+    return new(item) if item.is_a?(File)
+  end
+
+  def initialize(file)
+    @file = file
+  end
+
+  def retrieve_into(destination)
+    @file.rewind
+    @file.binmode
+    IO.copy_stream(@file, destination)
+  ensure
+    @file.close
+  end
+
+  def restartable?(_exception)
+    false
+  end
+end

lib/zipline/retrievers/http_retriever.rb

@@ -0,0 +1,34 @@
+class Zipline::HTTPRetriever
+  def self.build_for(url_or_uri)
+    return unless url_or_uri.is_a?(URI) || url_or_uri.is_a?(String) && url_or_uri.start_with?("http")
+
+    uri = begin
+      URI.parse(url_or_uri)
+    rescue
+      return
+    end
+
+    return new(uri) if uri.is_a?(URI::HTTP) || uri.is_a?(URI::HTTPS)
+  end
+
+  def initialize(uri)
+    @uri = uri
+  end
+
+  def retrieve_into(destination)
+    Net::HTTP.get_response(@uri) do |response|
+      response.read_body do |chunk|
+        destination.write(chunk)
+      end
+    end
+  end
+
+  def restartable?(exception)
+    restartables = [
+      Net::HTTPServerError,
+      Net::HTTPClientException,
+      Net::HTTPServiceUnavailable
+    ]
+    restartables.any? { |net_http_server_error_class| net_http_server_error_class === exception }
+  end
+end

lib/zipline/retrievers/io_retriever.rb

@@ -0,0 +1,17 @@
+class Zipline::IORetriever
+  def self.build_for(item)
+    return new(item) if item.respond_to?(:read) && item.respond_to?(:read_nonblock)
+  end
+
+  def initialize(io)
+    @io = io
+  end
+
+  def retrieve_into(destination)
+    IO.copy_stream(@io, destination)
+  end
+
+  def restartable?(_exception)
+    false
+  end
+end

lib/zipline/retrievers/paperclip_retriever.rb

@@ -0,0 +1,10 @@
+class Zipline::PaperclipRetriever
+  def self.build_for(item)
+    return unless defined?(Paperclip) && item.is_a?(Paperclip::Attachment)
+    if item.options[:storage] == :filesystem
+      Zipline::FileRetriever.build_for(File.open(item.path, "rb"))
+    else
+      Zipline::HTTPRetriever.build_for(file.expiring_url)
+    end
+  end
+end

lib/zipline/retrievers/shrine_retriever.rb

@@ -0,0 +1,18 @@
+class Zipline::ShrineRetriever
+  def self.build_for(item)
+    return unless defined?(Shrine::UploadedFile) && item.is_a?(Shrine::UploadedFile)
+    new(item)
+  end
+
+  def initialize(shrine_uploaded_file)
+    @shrine_uploaded_file = shrine_uploaded_file
+  end
+
+  def retrieve_into(destination)
+    @shrine_uploaded_file.stream(destination)
+  end
+
+  def restartable?(_exception)
+    false
+  end
+end

lib/zipline/retrievers/string_retriever.rb

@@ -0,0 +1,25 @@
+class Zipline::StringRetriever
+  def self.build_for(item)
+    return unless item.is_a?(String)
+    new(item)
+  end
+
+  def initialize(string)
+    @string = string
+  end
+
+  def retrieve_into(destination)
+    chunk_size = 1024
+    offset = 0
+    loop do
+      bytes = @string.byteslice(offset, chunk_size)
+      offset += chunk_size
+      destination.write(bytes)
+      break if bytes.nil?
+    end
+  end
+
+  def restartable?(_exception)
+    false
+  end
+end

lib/zipline/zip_handler.rb

@@ -1,98 +1,60 @@
-module Zipline
-  class ZipHandler
-    # takes an array of pairs [[uploader, filename], ... ]
-    def initialize(streamer, logger)
-      @streamer = streamer
-      @logger = logger
-    end
-
-    def handle_file(file, name, options)
-      normalized_file = normalize(file)
-      write_file(normalized_file, name, options)
-    rescue => e
-      # Since most APM packages do not trace errors occurring within streaming
-      # Rack bodies, it can be helpful to print the error to the Rails log at least
-      error_message = "zipline: an exception (#{e.inspect}) was raised  when serving the ZIP body."
-      error_message += " The error occurred when handling file #{name.inspect}"
-      @logger.error(error_message) if @logger
-      raise
-    end
+class Zipline::ZipHandler
+  def initialize(streamer, logger)
+    @streamer = streamer
+    @logger = logger
+  end
 
-    # This extracts either a url or a local file from the provided file.
-    # Currently support carrierwave and paperclip local and remote storage.
-    # returns a hash of the form {url: aUrl} or {file: anIoObject}
-    def normalize(file)
-      if defined?(CarrierWave::Uploader::Base) && file.is_a?(CarrierWave::Uploader::Base)
-        file = file.file
-      end
+  def handle_file(file, name, options)
+    write_item(file, name, options)
+  rescue => e
+    # Since most APM packages do not trace errors occurring within streaming
+    # Rack bodies, it can be helpful to print the error to the Rails log at least
+    error_message = "zipline: an exception (#{e.inspect}) was raised  when serving the ZIP body."
+    error_message += " The error occurred when handling file #{name.inspect} which was a #{file.class}"
+    @logger&.error(error_message)
+    raise
+  end
 
-      if defined?(Paperclip) && file.is_a?(Paperclip::Attachment)
-        if file.options[:storage] == :filesystem
-          {file: File.open(file.path)}
-        else
-          {url: file.expiring_url}
+  def write_item(item, name, options)
+    retriever = pick_retriever_for(item)
+    attempts = 0
+    begin
+      attempts += 1
+      @streamer.write_file(name, **options.slice(:modification_time)) do |writable|
+        ActiveSupport::Notifications.instrument("zipline.retrieve_and_write", {retriever_class: retriever.class.to_s, filename: name}) do
+          retriever.retrieve_into(writable)
         end
-      elsif defined?(CarrierWave::Storage::Fog::File) && file.is_a?(CarrierWave::Storage::Fog::File)
-        {url: file.url}
-      elsif defined?(CarrierWave::SanitizedFile) && file.is_a?(CarrierWave::SanitizedFile)
-        {file: File.open(file.path)}
-      elsif is_io?(file)
-        {file: file}
-      elsif defined?(ActiveStorage::Blob) && file.is_a?(ActiveStorage::Blob)
-        {blob: file}
-      elsif is_active_storage_attachment?(file) || is_active_storage_one?(file)
-        {blob: file.blob}
-      elsif file.respond_to? :url
-        {url: file.url}
-      elsif file.respond_to? :path
-        {file: File.open(file.path)}
-      elsif file.respond_to? :file
-        {file: File.open(file.file)}
-      elsif is_url?(file)
-        {url: file}
-      else
-        raise(ArgumentError, 'Bad File/Stream')
       end
-    end
-
-    def write_file(file, name, options)
-      @streamer.write_file(name, **options.slice(:modification_time)) do |writer_for_file|
-        if file[:url]
-          the_remote_uri = URI(file[:url])
-
-          Net::HTTP.get_response(the_remote_uri) do |response|
-            response.read_body do |chunk|
-              writer_for_file << chunk
-            end
-          end
-        elsif file[:file]
-          IO.copy_stream(file[:file], writer_for_file)
-          file[:file].close
-        elsif file[:blob]
-          file[:blob].download { |chunk| writer_for_file << chunk }
-        else
-          raise(ArgumentError, 'Bad File/Stream')
-        end
+    rescue => exception
+      # If an exception is raised and it is known to be caused by the data retrieval from
+      # remote, we can retry outputting this particular file. ZipKit will rollback the file
+      # before raising the exception to us, so we can just redo the `write_file` call.
+      if retriever.restartable?(exception)
+        @logger.warn { "Reattempting of #{name.inspect} will be reattempted after #{exception}" }
+        retry
+      else
+        @logger.warn { "Retrieval of #{name.inspect} cannot be restarted after #{exception}, abprting" }
+        raise
       end
     end
+  end
 
-    private
-
-    def is_io?(io_ish)
-      io_ish.respond_to? :read
-    end
-
-    def is_active_storage_attachment?(file)
-      defined?(ActiveStorage::Attachment) && file.is_a?(ActiveStorage::Attachment)
-    end
-
-    def is_active_storage_one?(file)
-      defined?(ActiveStorage::Attached::One) && file.is_a?(ActiveStorage::Attached::One)
-    end
-
-    def is_url?(url)
-      url = URI.parse(url) rescue false
-      url.kind_of?(URI::HTTP) || url.kind_of?(URI::HTTPS)
-    end
+  def pick_retriever_for(item)
+    retriever_classes = [
+      Zipline::CarrierwaveRetriever,
+      Zipline::ActiveStorageRetriever,
+      Zipline::PaperclipRetriever,
+      Zipline::ShrineRetriever,
+      Zipline::FileRetriever,
+      Zipline::IORetriever,
+      Zipline::HTTPRetriever,
+      Zipline::StringRetriever
+    ]
+    retriever_classes.each do |retriever_class|
+      maybe_retriever = retriever_class.build_for(item)
+      return maybe_retriever if maybe_retriever
+    end
+
+    raise "Don't know how to handle a file in the shape of #{item.inspect}"
   end
 end