1 files changed, 117 insertions, 0 deletions
diff --git a/lib/gitlab/import_export/remote_stream_upload.rb b/lib/gitlab/import_export/remote_stream_upload.rb
new file mode 100644
index 00000000000..f3bd241c0bd
--- /dev/null
+++ b/lib/gitlab/import_export/remote_stream_upload.rb
@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+
+# This class downloads a file from one URL and uploads it to another URL
+# without having to save the file on the disk and loading the whole file in
+# memory. The download and upload are performed in chunks size of
+# `buffer_size`. A chunk is downloaded, then uploaded, then a next chunk is
+# downloaded and uploaded. This repeats until all the file is processed.
+
+module Gitlab
+  module ImportExport
+    class RemoteStreamUpload
+      def initialize(download_url:, upload_url:, options: {})
+        @download_url = download_url
+        @upload_url = upload_url
+        @upload_method = options[:upload_method] || :post
+        @upload_content_type = options[:upload_content_type] || 'application/gzip'
+      end
+
+      def execute
+        receive_data(download_url) do |response, chunks|
+          send_data(upload_url, response.content_length, chunks) do |response|
+            if response.code != '200'
+              raise StreamError.new("Invalid response code while uploading file. Code: #{response.code}", response.body)
+            end
+          end
+        end
+      end
+      class StreamError < StandardError
+        attr_reader :response_body
+
+        def initialize(message, response_body = '')
+          super(message)
+          @response_body = response_body
+        end
+      end
+      class ChunkStream
+        DEFAULT_BUFFER_SIZE = 128.kilobytes
+
+        def initialize(chunks)
+          @chunks = chunks
+          @last_chunk = nil
+          @end_of_chunks = false
+        end
+
+        def read(n1 = nil, n2 = nil)
+          ensure_chunk&.read(n1, n2)
+        end
+
+        private
+
+        def ensure_chunk
+          return @last_chunk if @last_chunk && !@last_chunk.eof?
+          return if @end_of_chunks
+
+          @last_chunk = read_next_chunk
+        end
+
+        def read_next_chunk
+          next_chunk = StringIO.new
+
+          begin
+            next_chunk.write(@chunks.next) until next_chunk.size > DEFAULT_BUFFER_SIZE
+          rescue StopIteration
+            @end_of_chunks = true
+          end
+
+          next_chunk.rewind
+
+          next_chunk
+        end
+      end
+
+      private
+
+      attr_reader :download_url, :upload_url, :upload_method, :upload_content_type, :logger
+
+      def receive_data(uri)
+        http = Gitlab::HTTPConnectionAdapter.new(URI(uri), {}).connection
+
+        http.start do
+          request = Net::HTTP::Get.new(uri)
+          http.request(request) do |response|
+            if response.code == '200'
+              yield(response, response.enum_for(:read_body))
+            else
+              raise StreamError.new(
+                "Invalid response code while downloading file. Code: #{response.code}",
+                response.body
+              )
+            end
+          end
+        end
+      end
+
+      def send_data(uri, content_length, chunks)
+        http = Gitlab::HTTPConnectionAdapter.new(URI(uri), {}).connection
+
+        http.start do
+          request = upload_request_class(upload_method).new(uri)
+          request.body_stream = ChunkStream.new(chunks)
+          request.content_length = content_length
+          request.content_type = upload_content_type
+
+          http.request(request) do |response|
+            yield(response)
+          end
+        end
+      end
+
+      def upload_request_class(upload_method)
+        return Net::HTTP::Put if upload_method == :put
+
+        Net::HTTP::Post
+      end
+    end
+  end
+end