1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
|
# 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
|
