s3_direct_multipart_upload 0.1.0

data/app/models/s3_direct_multipart_upload/upload_session.rb

@@ -0,0 +1,394 @@
+# == Schema Information
+#
+# Table name: s3_direct_multipart_upload_sessions
+#
+#  id           :bigint           not null, primary key
+#  bucket       :string(255)      not null
+#  byte_size    :bigint           not null
+#  chunk_size   :integer
+#  content_type :string(255)
+#  expires_at   :datetime         indexed
+#  filename     :string(255)      not null
+#  key_prefix   :string(255)      not null
+#  metadata     :json             not null
+#  status       :string(255)      default("pending"), not null
+#  created_at   :datetime         not null
+#  updated_at   :datetime         not null
+#  session_id   :string(255)      not null, uniquely indexed
+#  upload_id    :string(255)      not null, uniquely indexed
+#
+require "digest"
+require "erb"
+require "uri"
+require "yaml"
+
+module S3DirectMultipartUpload
+  class UploadSession < ::ApplicationRecord
+    self.table_name = "s3_direct_multipart_upload_sessions"
+
+    MAX_PARTS = 10_000
+    MIN_MULTIPART_CHUNK_SIZE = 5.megabytes
+    MAX_CHUNK_SIZE = 5.gigabytes
+    SINGLE_PART_THRESHOLD = MIN_MULTIPART_CHUNK_SIZE
+    PRESIGNED_URL_TTL = 15.minutes
+
+    has_many :parts,
+             class_name: "S3DirectMultipartUpload::UploadSessionPart",
+             inverse_of: :upload_session,
+             dependent: :destroy
+
+    enum :status, {
+      pending: "pending",
+      uploading: "uploading",
+      completed: "completed",
+      aborted: "aborted"
+    }
+
+    validates :session_id, :upload_id, :bucket, :key_prefix, :status, presence: true
+    validates :session_id, :upload_id, uniqueness: true
+    validates :byte_size, numericality: { greater_than: 0 }
+    validates :chunk_size, numericality: { greater_than: 0 }, allow_nil: true
+    validate :validate_chunk_size_bounds, :validate_expected_parts
+
+    before_validation :assign_defaults, on: :create
+
+    def self.start!(filename:, byte_size:, content_type:, chunk_size: nil, metadata: {})
+      byte_size = byte_size.to_i
+      raise ArgumentError, "byte_size must be positive" if byte_size <= 0
+
+      resolved_chunk_size = resolve_chunk_size(byte_size:, chunk_size:)
+      expected_parts = expected_parts_for(byte_size:, chunk_size: resolved_chunk_size)
+      raise ArgumentError, "part number limit exceeded" if expected_parts > MAX_PARTS
+
+      session_id = SecureRandom.uuid
+      key_prefix = default_key_prefix(session_id)
+      object_key = build_object_key(key_prefix:, filename:)
+
+      response = s3_client.create_multipart_upload(
+        bucket: default_bucket,
+        key: object_key,
+        content_type:
+      )
+
+      create!(
+        session_id:,
+        upload_id: response.upload_id,
+        bucket: response.bucket || default_bucket,
+        key_prefix:,
+        filename:,
+        byte_size:,
+        content_type:,
+        chunk_size: resolved_chunk_size,
+        metadata: metadata
+      ).tap(&:schedule_expiration!)
+    end
+
+    def presigned_parts(part_numbers:)
+      raise ArgumentError, "part number limit exceeded" if part_numbers.size > MAX_PARTS
+
+      part_numbers.map do |number|
+        disk_service? ? dev_presigned_part(number) : aws_presigned_part(number)
+      end
+    end
+
+    def report_part!(part_number:, etag:)
+      part_number = part_number.to_i
+      raise ArgumentError, "part_number is invalid" if part_number.zero? || part_number.negative?
+      raise ArgumentError, "part_number exceeds expected parts" if expected_parts.positive? && part_number > expected_parts
+      raise ArgumentError, "upload already completed or aborted" if completed? || aborted?
+
+      transaction do
+        record = parts.find_or_initialize_by(part_number:)
+        record.update!(etag:, uploaded_at: Time.current)
+        uploading! if pending?
+      end
+    end
+
+    def complete!(checksum: nil)
+      missing = missing_parts
+      raise "missing parts: #{missing}" if missing.present?
+
+      self.class.s3_client.complete_multipart_upload(
+        bucket:,
+        key: object_key,
+        upload_id:,
+        multipart_upload: {
+          parts: parts.order(:part_number).map { { part_number: _1.part_number, etag: _1.etag } }
+        }
+      )
+      completed!
+      update!(metadata: metadata.merge("checksum" => checksum).compact)
+      combine_dev_parts_if_needed
+      self
+    end
+
+    def object_key
+      @object_key ||= self.class.build_object_key(key_prefix:, filename:)
+    end
+
+    def parts_uploaded?
+      expected_parts.positive? && missing_parts.empty?
+    end
+
+    def expected_parts
+      return 0 if chunk_size.blank? || byte_size.blank?
+
+      self.class.expected_parts_for(byte_size:, chunk_size:)
+    end
+
+    def part_limit
+      MAX_PARTS
+    end
+
+    def missing_parts
+      return [] unless expected_parts.positive?
+
+      uploaded = parts.pluck(:part_number).uniq.sort
+      expected = (1..expected_parts).to_a
+      expected - uploaded
+    end
+
+    def schedule_expiration!
+      update!(expires_at: 24.hours.from_now)
+    end
+
+    def default_chunk_size
+      self.class.default_chunk_size
+    end
+
+    private
+
+    def assign_defaults
+      self.session_id ||= SecureRandom.uuid
+      self.status ||= :pending
+      self.bucket ||= self.class.default_bucket
+      self.chunk_size ||= self.class.default_chunk_size
+      self.metadata ||= {}
+      self.key_prefix ||= self.class.default_key_prefix(session_id)
+    end
+
+    def presigner
+      @presigner ||= Aws::S3::Presigner.new(client: self.class.s3_client)
+    end
+
+    def disk_service?
+      self.class.disk_service?
+    end
+
+    def aws_presigned_part(number)
+      {
+        part_number: number,
+        url: presigner.presigned_url(
+          :upload_part,
+          bucket:,
+          key: object_key,
+          upload_id:,
+          part_number: number,
+          expires_in: PRESIGNED_URL_TTL.to_i
+        ),
+        expires_in: PRESIGNED_URL_TTL.to_i
+      }
+    end
+
+    def dev_presigned_part(number)
+      {
+        part_number: number,
+        url: self.class.dev_presigned_url(object_key:, upload_id:, part_number: number),
+        expires_in: PRESIGNED_URL_TTL.to_i
+      }
+    end
+
+    def validate_chunk_size_bounds
+      return if chunk_size.blank? || byte_size.blank?
+
+      if single_part?
+        errors.add(:chunk_size, "must equal byte_size for single part") unless chunk_size == byte_size
+        return
+      end
+
+      errors.add(:chunk_size, "must be at least #{MIN_MULTIPART_CHUNK_SIZE}") if chunk_size < MIN_MULTIPART_CHUNK_SIZE
+      errors.add(:chunk_size, "must be at most #{MAX_CHUNK_SIZE}") if chunk_size > MAX_CHUNK_SIZE
+    end
+
+    def validate_expected_parts
+      return if chunk_size.blank? || byte_size.blank?
+
+      errors.add(:base, "part number limit exceeded") if expected_parts > MAX_PARTS
+    end
+
+    def single_part?
+      byte_size.present? && byte_size < SINGLE_PART_THRESHOLD
+    end
+
+    def combine_dev_parts_if_needed
+      return unless self.class.dev_mode_enabled?
+      return unless self.class.disk_service?
+
+      S3DirectMultipartUpload::Dev::Storage.combine_parts(self)
+    rescue StandardError => e
+      Rails.logger.warn("[S3DirectMultipartUpload] failed to combine parts in dev mode: #{e.message}")
+    end
+
+    class << self
+      def default_bucket
+        storage_config.fetch(:bucket)
+      end
+
+      def default_key_prefix(session_id)
+        date_prefix = Time.current.strftime("%Y%m%d")
+        File.join(upload_namespace, date_prefix, session_id)
+      end
+
+      def upload_namespace
+        "uploads/s3_direct"
+      end
+
+      def default_chunk_size
+        128.megabytes
+      end
+
+      def resolve_chunk_size(byte_size:, chunk_size:)
+        return byte_size if byte_size < SINGLE_PART_THRESHOLD
+
+        requested = chunk_size.presence&.to_i || default_chunk_size
+        raise ArgumentError, "chunk_size must be positive" if requested <= 0
+        raise ArgumentError, "chunk_size must be at least #{MIN_MULTIPART_CHUNK_SIZE}" if requested < MIN_MULTIPART_CHUNK_SIZE
+        raise ArgumentError, "chunk_size must be at most #{MAX_CHUNK_SIZE}" if requested > MAX_CHUNK_SIZE
+
+        requested
+      end
+
+      def expected_parts_for(byte_size:, chunk_size:)
+        (byte_size.to_f / chunk_size).ceil
+      end
+
+      def s3_client
+        @s3_client ||= begin
+          config = storage_config
+          client_options = {
+            region: config.fetch(:region),
+            endpoint: config[:endpoint],
+            force_path_style: config[:force_path_style],
+            stub_responses: stub_responses_option(config)
+          }.compact
+          Aws::S3::Client.new(**client_options)
+        end
+      end
+
+      def build_object_key(key_prefix:, filename:)
+        File.join(key_prefix, filename)
+      end
+
+      def storage_config
+        @storage_config ||= begin
+          config = storage_configurations["s3_direct_multipart"] || storage_configurations[:s3_direct_multipart]
+          raise configuration_error("ActiveStorage service `s3_direct_multipart` is not configured") unless config
+
+          config = normalize_config(config.deep_symbolize_keys)
+          validate_config!(config)
+          config
+        end
+      end
+
+      def storage_configurations
+        explicit = Rails.application.config.respond_to?(:active_storage) &&
+                   Rails.application.config.active_storage.respond_to?(:service_configurations) &&
+                   Rails.application.config.active_storage.service_configurations.present?
+        return Rails.application.config.active_storage.service_configurations if explicit
+
+        path = storage_config_path
+        raise configuration_error("ActiveStorage service configurations are empty; add config/storage.yml or config/storage/<env>.yml") unless path
+
+        erb = ERB.new(path.read).result
+        YAML.safe_load(erb, aliases: true) || {}
+      end
+
+      def storage_config_path
+        env_path = Rails.root.join("config/storage/#{Rails.env}.yml")
+        return env_path if env_path.exist?
+
+        default_path = Rails.root.join("config/storage.yml")
+        default_path if default_path.exist?
+      end
+
+      def validate_config!(config)
+        service = config[:service].to_s.downcase
+        raise configuration_error("ActiveStorage service `s3_direct_multipart` must set service: S3 or Disk") unless %w[s3 disk].include?(service)
+        raise configuration_error("ActiveStorage service `s3_direct_multipart` requires bucket") if config[:bucket].blank?
+        raise configuration_error("ActiveStorage service `s3_direct_multipart` requires region") if config[:region].blank?
+      end
+
+      def normalize_config(config)
+        service = config[:service].to_s.downcase
+        return config if service == "s3"
+
+        config.merge(
+          service: "Disk",
+          bucket: config[:bucket] || "s3_direct_disk",
+          region: config[:region] || "us-east-1",
+          stub_responses: true
+        )
+      end
+
+      def disk_service?
+        storage_config[:service].to_s.casecmp("disk").zero?
+      end
+
+      def dev_presigned_url(object_key:, upload_id:, part_number:)
+        raise configuration_error("dev presigned URLs are only available for Disk service") unless disk_service?
+
+        expires_at = PRESIGNED_URL_TTL.from_now
+        signature = dev_signature(
+          path: object_key,
+          expires_at:,
+          upload_id:,
+          part_number:
+        )
+
+        path = S3DirectMultipartUpload::Engine.routes.url_helpers.dev_storage_upload_path(path: object_key)
+        url = URI.join(dev_endpoint_base_url, path)
+        query = Rack::Utils.build_query(
+          uploadId: upload_id,
+          partNumber: part_number,
+          expires_at: expires_at.to_i,
+          signature:
+        )
+        url.query = query
+        url.to_s
+      end
+
+      def dev_signature(path:, expires_at:, upload_id:, part_number:)
+        secret = ENV.fetch("S3_DIRECT_DEV_SECRET", "dev-secret")
+        normalized_path = path.to_s.sub(/\A\//, "")
+        normalized_part = part_number.presence&.to_s
+        Digest::SHA256.hexdigest([ normalized_path, expires_at.to_i, upload_id.to_s, normalized_part, secret ].compact.join(":"))
+      end
+
+      def dev_endpoint_base_url
+        ENV.fetch("S3_DIRECT_DEV_ENDPOINT", "http://localhost:3000")
+      end
+
+      def stub_responses_option(config)
+        return config[:stub_responses] unless disk_service? && config[:stub_responses]
+        return config[:stub_responses] if config[:stub_responses].respond_to?(:to_hash)
+
+        {
+          create_multipart_upload: ->(*) {
+            {
+              upload_id: SecureRandom.uuid,
+              bucket: config[:bucket] || default_bucket
+            }
+          }
+        }
+      end
+
+      def dev_mode_enabled?
+        Rails.env.development? || Rails.env.test?
+      end
+
+      def configuration_error(message)
+        RuntimeError.new(message)
+      end
+    end
+  end
+end

data/app/models/s3_direct_multipart_upload/upload_session_part.rb

@@ -0,0 +1,34 @@
+# == Schema Information
+#
+# Table name: s3_direct_multipart_upload_session_parts
+#
+#  id                :bigint           not null, primary key
+#  etag              :string(255)      not null
+#  part_number       :integer          not null, uniquely indexed => [upload_session_id]
+#  uploaded_at       :datetime         not null
+#  upload_session_id :bigint           not null, indexed, uniquely indexed => [part_number]
+#
+# Foreign Keys
+#
+#  fk_rails_...  (upload_session_id => s3_direct_multipart_upload_sessions.id)
+#
+module S3DirectMultipartUpload
+  class UploadSessionPart < ::ApplicationRecord
+    self.table_name = "s3_direct_multipart_upload_session_parts"
+
+    belongs_to :upload_session,
+               class_name: "S3DirectMultipartUpload::UploadSession",
+               inverse_of: :parts
+
+    validates :part_number, presence: true, numericality: { greater_than: 0 }
+    validates :etag, presence: true
+
+    after_create :touch_upload_session
+
+    private
+
+    def touch_upload_session
+      upload_session.touch
+    end
+  end
+end

data/app/services/s3_direct_multipart_upload/dev/storage.rb

@@ -0,0 +1,85 @@
+require "fileutils"
+require "stringio"
+
+module S3DirectMultipartUpload
+  module Dev
+    module Storage
+      module_function
+
+      def root
+        Rails.root.join("tmp/s3_direct_multipart_upload")
+      end
+
+      def combined_path(upload_session)
+        root.join(upload_session.object_key.to_s.sub(/\A\//, ""))
+      end
+
+      def parts_dir(upload_session)
+        combined_path(upload_session).dirname.join("#{combined_path(upload_session).basename}.parts")
+      end
+
+      def combine_parts(upload_session)
+        raise "download is available only in development/test" unless dev_environment?
+        ensure_directories(upload_session)
+
+        if disk_service?(upload_session)
+          combine_disk_parts(upload_session)
+        else
+          download_s3_object(upload_session)
+        end
+      end
+
+      def combined_file_exist?(upload_session)
+        combined_path(upload_session).exist?
+      end
+
+      def dev_environment?
+        Rails.env.development? || Rails.env.test?
+      end
+
+      def disk_service?(upload_session)
+        upload_session.class.send(:disk_service?)
+      end
+
+      def ensure_directories(upload_session)
+        FileUtils.mkdir_p(combined_path(upload_session).dirname)
+        FileUtils.mkdir_p(parts_dir(upload_session))
+      end
+
+      def combine_disk_parts(upload_session)
+        expected = upload_session.expected_parts
+        raise "upload is not completed" unless upload_session.completed?
+
+        if expected <= 1
+          single_part_path = parts_dir(upload_session).join("1")
+          source = single_part_path.exist? ? single_part_path : combined_path(upload_session)
+          raise "part is missing" unless source.exist?
+
+          FileUtils.cp(source, combined_path(upload_session))
+          return combined_path(upload_session)
+        end
+
+        io = StringIO.new
+        (1..expected).each do |part_number|
+          part_path = parts_dir(upload_session).join(part_number.to_s)
+          raise "missing part #{part_number}" unless part_path.exist?
+
+          io.write(File.binread(part_path))
+        end
+
+        File.binwrite(combined_path(upload_session), io.string)
+        combined_path(upload_session)
+      end
+
+      def download_s3_object(upload_session)
+        client = upload_session.class.s3_client
+        client.get_object(
+          bucket: upload_session.bucket,
+          key: upload_session.object_key,
+          response_target: combined_path(upload_session)
+        )
+        combined_path(upload_session)
+      end
+    end
+  end
+end

data/config/initializers/s3_direct_dev.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+if Rails.env.development? || Rails.env.test?
+  Aws.config[:s3] ||= {}
+  Aws.config[:s3][:stub_responses] = {
+    create_multipart_upload: ->(*) {
+      {
+        upload_id: SecureRandom.uuid,
+        bucket: "dev-bucket"
+      }
+    }
+  }
+end

data/config/routes.rb

@@ -0,0 +1,15 @@
+S3DirectMultipartUpload::Engine.routes.draw do
+  post "s3_direct_multipart_uploads", to: "upload_sessions#create", as: :s3_direct_multipart_uploads
+  post "s3_direct_multipart_uploads/:upload_session_id/parts", to: "upload_parts#create", as: :s3_direct_multipart_upload_parts
+  post "s3_direct_multipart_uploads/:upload_session_id/complete", to: "upload_completions#create", as: :s3_direct_multipart_upload_complete
+
+  if Rails.env.development? || Rails.env.test?
+    put "s3_direct_multipart_upload/dev/storage/*path",
+        to: "dev/storage#upload",
+        as: :dev_storage_upload,
+        format: false
+    get "s3_direct_multipart_upload/dev/storage/:upload_session_id/download",
+        to: "dev/storage_downloads#show",
+        as: :dev_storage_download
+  end
+end

data/config/storage/development.yml

@@ -0,0 +1,3 @@
+s3_direct_multipart:
+  service: Disk
+  root: <%= Rails.root.join('tmp/storage') %>

data/config/storage/test.yml

@@ -0,0 +1,3 @@
+s3_direct_multipart:
+  service: Disk
+  root: <%= Rails.root.join('tmp/storage') %>

data/db/migrate/20251113090000_create_s3_direct_multipart_upload_tables.rb

@@ -0,0 +1,32 @@
+class CreateS3DirectMultipartUploadTables < ActiveRecord::Migration[7.1]
+  def change
+    create_table :s3_direct_multipart_upload_sessions do |t|
+      t.string  :session_id, null: false
+      t.string  :upload_id,  null: false
+      t.string  :bucket,     null: false
+      t.string  :key_prefix, null: false
+      t.string  :status,     null: false, default: 'pending'
+      t.bigint  :byte_size,  null: false
+      t.string  :filename,   null: false
+      t.string  :content_type
+      t.integer :chunk_size
+      t.datetime :expires_at
+      t.json :metadata, null: false
+
+      t.timestamps
+
+      t.index :session_id, unique: true
+      t.index :upload_id, unique: true
+      t.index :expires_at
+    end
+
+    create_table :s3_direct_multipart_upload_session_parts do |t|
+      t.references :upload_session, null: false, foreign_key: { to_table: :s3_direct_multipart_upload_sessions }
+      t.integer    :part_number,    null: false
+      t.string     :etag,           null: false
+      t.datetime   :uploaded_at,    null: false
+
+      t.index [ :upload_session_id, :part_number ], unique: true, name: 'index_s3_direct_parts_on_session_and_number'
+    end
+  end
+end