Skip to content

Commit 61ec599

Browse files
johnnyshieldspicandocodigo
johnnyshields
authored and
picandocodigo
committed
Fixes compression on elasticsearch-transport
- "compression" option should compress outbound requests. - Decompress compressed responses, whether or not use_compression is set. - Properly set header on gzipped request body - Perform compression inside HTTP adapters - Explicitly require zlib library - Content-Encoding is set in the HTTP adapters, so we don't need it in the apply_headers method - Add specs for curb, faraday, and manticore - Update curb_spec.rb
1 parent 4cf4719 commit 61ec599

File tree

9 files changed

+459
-20
lines changed

9 files changed

+459
-20
lines changed

‎elasticsearch-transport/lib/elasticsearch/transport.rb

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
require 'uri'
1919
require 'time'
2020
require 'timeout'
21+
require 'zlib'
2122
require 'multi_json'
2223
require 'faraday'
2324

‎elasticsearch-transport/lib/elasticsearch/transport/transport/base.rb

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -224,7 +224,7 @@ def __raise_transport_error(response)
224224
# @api private
225225
#
226226
def __convert_to_json(o=nil, options={})
227-
o = o.is_a?(String) ? o : serializer.dump(o, options)
227+
o.is_a?(String) ? o : serializer.dump(o, options)
228228
end
229229

230230
# Returns a full URL based on information from host
@@ -369,17 +369,38 @@ def host_unreachable_exceptions
369369

370370
USER_AGENT_STR = 'User-Agent'.freeze
371371
USER_AGENT_REGEX = /user\-?\_?agent/
372+
ACCEPT_ENCODING = 'Accept-Encoding'.freeze
373+
CONTENT_ENCODING = 'Content-Encoding'.freeze
372374
CONTENT_TYPE_STR = 'Content-Type'.freeze
373375
CONTENT_TYPE_REGEX = /content\-?\_?type/
374376
DEFAULT_CONTENT_TYPE = 'application/json'.freeze
375377
GZIP = 'gzip'.freeze
376-
ACCEPT_ENCODING = 'Accept-Encoding'.freeze
377378
GZIP_FIRST_TWO_BYTES = '1f8b'.freeze
378379
HEX_STRING_DIRECTIVE = 'H*'.freeze
379380
RUBY_ENCODING = '1.9'.respond_to?(:force_encoding)
380381

382+
def compress_request(body, headers)
383+
if body
384+
headers ||= {}
385+
386+
if gzipped?(body)
387+
headers[CONTENT_ENCODING] = GZIP
388+
elsif use_compression?
389+
headers[CONTENT_ENCODING] = GZIP
390+
gzip = Zlib::GzipWriter.new(StringIO.new)
391+
gzip << body
392+
body = gzip.close.string
393+
else
394+
headers.delete(CONTENT_ENCODING)
395+
end
396+
elsif headers
397+
headers.delete(CONTENT_ENCODING)
398+
end
399+
400+
[body, headers]
401+
end
402+
381403
def decompress_response(body)
382-
return body unless use_compression?
383404
return body unless gzipped?(body)
384405

385406
io = StringIO.new(body)
@@ -392,6 +413,7 @@ def decompress_response(body)
392413
end
393414

394415
def gzipped?(body)
416+
return unless body
395417
body[0..1].unpack(HEX_STRING_DIRECTIVE)[0] == GZIP_FIRST_TWO_BYTES
396418
end
397419

‎elasticsearch-transport/lib/elasticsearch/transport/transport/http/curb.rb

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,29 +19,31 @@ module Elasticsearch
1919
module Transport
2020
module Transport
2121
module HTTP
22-
2322
# Alternative HTTP transport implementation, using the [_Curb_](https://rubygems.org/gems/curb) client.
2423
#
2524
# @see Transport::Base
2625
#
2726
class Curb
2827
include Base
29-
3028
# Performs the request by invoking {Transport::Base#perform_request} with a block.
3129
#
3230
# @return [Response]
3331
# @see Transport::Base#perform_request
3432
#
3533
def perform_request(method, path, params={}, body=nil, headers=nil, opts={})
36-
super do |connection, url|
34+
super do |connection, _url|
3735
connection.connection.url = connection.full_url(path, params)
36+
body = body ? __convert_to_json(body) : nil
37+
body, headers = compress_request(body, headers)
3838

3939
case method
4040
when 'HEAD'
4141
connection.connection.set :nobody, true
4242
when 'GET', 'POST', 'PUT', 'DELETE'
4343
connection.connection.set :nobody, false
44-
connection.connection.put_data = __convert_to_json(body) if body
44+
45+
connection.connection.put_data = body if body
46+
4547
if headers
4648
if connection.connection.headers
4749
connection.connection.headers.merge!(headers)

‎elasticsearch-transport/lib/elasticsearch/transport/transport/http/faraday.rb

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -44,12 +44,13 @@ def perform_request(method, path, params = {}, body = nil, headers = nil, opts =
4444
headers
4545
end
4646

47-
response = connection.connection.run_request(
48-
method.downcase.to_sym,
49-
url,
50-
(body ? __convert_to_json(body) : nil),
51-
headers
52-
)
47+
body = body ? __convert_to_json(body) : nil
48+
body, headers = compress_request(body, headers)
49+
50+
response = connection.connection.run_request(method.downcase.to_sym,
51+
url,
52+
body,
53+
headers)
5354

5455
Response.new response.status, decompress_response(response.body), response.headers
5556
end

‎elasticsearch-transport/lib/elasticsearch/transport/transport/http/manticore.rb

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,10 @@ def build_client(options={})
8383
#
8484
def perform_request(method, path, params={}, body=nil, headers=nil, opts={})
8585
super do |connection, url|
86-
params[:body] = __convert_to_json(body) if body
86+
body = body ? __convert_to_json(body) : nil
87+
body, headers = compress_request(body, @request_options[:headers])
88+
89+
params[:body] = body if body
8790
params[:headers] = headers if headers
8891
params = params.merge @request_options
8992
case method

‎elasticsearch-transport/spec/elasticsearch/transport/client_spec.rb

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1741,7 +1741,7 @@
17411741
end
17421742

17431743
it 'sets the Accept-Encoding header' do
1744-
expect(client.transport.connections[0].connection.headers['Accept-Encoding'])
1744+
expect(client.transport.connections[0].connection.headers['Accept-Encoding']).to eq 'gzip'
17451745
end
17461746

17471747
it 'preserves the other headers' do
@@ -1760,7 +1760,7 @@
17601760
end
17611761

17621762
it 'sets the Accept-Encoding header' do
1763-
expect(client.transport.connections[0].connection.headers['Accept-Encoding'])
1763+
expect(client.transport.connections[0].connection.headers['Accept-Encoding']).to eq 'gzip'
17641764
end
17651765

17661766
it 'preserves the other headers' do
@@ -1779,7 +1779,7 @@
17791779
end
17801780

17811781
it 'sets the Accept-Encoding header' do
1782-
expect(client.transport.connections[0].connection.headers['Accept-Encoding'])
1782+
expect(client.transport.connections[0].connection.headers['Accept-Encoding']).to eq 'gzip'
17831783
end
17841784

17851785
it 'preserves the other headers' do
@@ -1798,7 +1798,7 @@
17981798
end
17991799

18001800
it 'sets the Accept-Encoding header' do
1801-
expect(client.transport.connections[0].connection.headers['Accept-Encoding'])
1801+
expect(client.transport.connections[0].connection.headers['Accept-Encoding']).to eq 'gzip'
18021802
end
18031803

18041804
it 'preserves the other headers' do
@@ -1817,7 +1817,7 @@
18171817
end
18181818

18191819
it 'sets the Accept-Encoding header' do
1820-
expect(client.transport.connections[0].connection.headers['Accept-Encoding'])
1820+
expect(client.transport.connections[0].connection.headers['Accept-Encoding']).to eq 'gzip'
18211821
end
18221822

18231823
it 'preserves the other headers' do
@@ -1841,7 +1841,7 @@
18411841
end
18421842

18431843
it 'sets the Accept-Encoding header' do
1844-
expect(client.transport.connections[0].connection.headers['Accept-Encoding'])
1844+
expect(client.transport.connections[0].connection.headers['Accept-Encoding']).to eq 'gzip'
18451845
end
18461846

18471847
it 'preserves the other headers' do

‎elasticsearch-transport/spec/elasticsearch/transport/http/curb_spec.rb

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
# Licensed to Elasticsearch B.V. under one or more contributor
2+
# license agreements. See the NOTICE file distributed with
3+
# this work for additional information regarding copyright
4+
# ownership. Elasticsearch B.V. licenses this file to you under
5+
# the Apache License, Version 2.0 (the "License"); you may
6+
# not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
unless defined?(JRUBY_VERSION)
19+
require_relative '../../../spec_helper'
20+
21+
describe Elasticsearch::Transport::Transport::HTTP::Curb do
22+
let(:client) do
23+
Elasticsearch::Transport::Client.new(transport_class: described_class)
24+
end
25+
26+
describe '#perform_request' do
27+
subject(:perform_request) { client.perform_request(*args) }
28+
let(:args) do
29+
['POST', '/', {}, body, headers]
30+
end
31+
let(:body) { '{"foo":"bar"}' }
32+
let(:headers) { { 'Content-Type' => 'application/x-ndjson' } }
33+
34+
before do
35+
allow_any_instance_of(Curl::Easy).to receive(:http).and_return(true)
36+
end
37+
38+
it 'convert body to json' do
39+
expect(client.transport).to receive(:__convert_to_json).with(body)
40+
perform_request
41+
end
42+
43+
it 'call compress_request' do
44+
expect(client.transport).to receive(:compress_request).with(body, headers)
45+
perform_request
46+
end
47+
48+
it 'return response' do
49+
expect(perform_request).to be_kind_of(Elasticsearch::Transport::Transport::Response)
50+
end
51+
52+
it 'put body' do
53+
expect(client.transport.connections.first.connection).to receive('put_data=').with(body)
54+
perform_request
55+
end
56+
57+
context 'when body nil' do
58+
let(:body) { nil }
59+
60+
it 'convert body to json' do
61+
expect(client.transport).not_to receive(:__convert_to_json)
62+
perform_request
63+
end
64+
65+
it 'call compress_request' do
66+
expect(client.transport).to receive(:compress_request).with(body, headers)
67+
perform_request
68+
end
69+
70+
it 'put body' do
71+
expect(client.transport.connections.first.connection).not_to receive('put_data=')
72+
perform_request
73+
end
74+
end
75+
76+
context 'when body is hash' do
77+
let(:body) { { foo: 'bar' } }
78+
let(:body_string) { '{"foo":"bar"}' }
79+
80+
it 'convert body to json' do
81+
expect(client.transport).to receive(:__convert_to_json).with(body)
82+
perform_request
83+
end
84+
85+
it 'call compress_request' do
86+
expect(client.transport).to receive(:compress_request).with(body_string, headers)
87+
perform_request
88+
end
89+
90+
it 'put body' do
91+
expect(client.transport.connections.first.connection).to receive('put_data=').with(body_string)
92+
perform_request
93+
end
94+
end
95+
96+
context 'when compression enabled' do
97+
let(:client) do
98+
Elasticsearch::Transport::Client.new(transport_class: described_class, compression: true)
99+
end
100+
let(:body_string) { '{"foo":"bar"}' }
101+
let(:compressed_body) do
102+
gzip = Zlib::GzipWriter.new(StringIO.new)
103+
gzip << body_string
104+
gzip.close.string
105+
end
106+
107+
before { allow(client.transport).to receive(:decompress_response).and_return('') }
108+
109+
it 'put compressed body' do
110+
expect(client.transport.connections.first.connection).to receive('put_data=').with(compressed_body)
111+
perform_request
112+
end
113+
114+
it 'set Content-Encoding header' do
115+
perform_request
116+
expect(client.transport.connections.first.connection.headers).to include('Content-Encoding')
117+
end
118+
119+
it 'set Content-Encoding to gzip' do
120+
perform_request
121+
expect(client.transport.connections.first.connection.headers['Content-Encoding']).to eql('gzip')
122+
end
123+
end
124+
end
125+
end
126+
end

0 commit comments

Comments
 (0)