Gzip compression/inflation handling

CHANGELOG.md

@@ -8,6 +8,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ### Deprecated
 ### Removed
 ### Fixed
+- Fixed client compression behavior, specifically applying compression on requests, always inflating gzipped responses, and properly handling `compression: false` ([#338](https://github.com/opensearch-project/opensearch-ruby/pull/338))
 ### Security
 
 ## [4.0.0]

lib/opensearch/transport/client.rb

@@ -145,6 +145,7 @@ def initialize(arguments = {}, &block)
         @arguments[:retry_on_failure]   ||= false
         @arguments[:reload_on_failure]  ||= false
         @arguments[:randomize_hosts]    ||= false
+        @arguments[:compression]        ||= false
         @arguments[:transport_options]  ||= {}
         @arguments[:http]               ||= {}
         @options[:http]                 ||= {}

lib/opensearch/transport/transport/base.rb

@@ -64,7 +64,7 @@ def initialize(arguments = {}, &block)
           @options[:retry_on_status] ||= []
 
           @block       = block
-          @compression = !@options[:compression].nil?
+          @compression = !!@options[:compression]
           @connections = __build_connections
 
           @serializer  = options[:serializer] || (options[:serializer_class] ? options[:serializer_class].new(self) : DEFAULT_SERIALIZER_CLASS.new(self))
@@ -389,21 +389,15 @@ def host_unreachable_exceptions
         DEFAULT_CONTENT_TYPE = 'application/json'.freeze
         GZIP = 'gzip'.freeze
         ACCEPT_ENCODING = 'Accept-Encoding'.freeze
+        CONTENT_ENCODING = 'Content-Encoding'.freeze
         GZIP_FIRST_TWO_BYTES = '1f8b'.freeze
         HEX_STRING_DIRECTIVE = 'H*'.freeze
-        RUBY_ENCODING = '1.9'.respond_to?(:force_encoding)
 
         def decompress_response(body)
-          return body unless use_compression?
+          return nil unless body
           return body unless gzipped?(body)
 
-          io = StringIO.new(body)
-          gzip_reader = if RUBY_ENCODING
-                          Zlib::GzipReader.new(io, encoding: 'ASCII-8BIT')
-                        else
-                          Zlib::GzipReader.new(io)
-                        end
-          gzip_reader.read
+          Zlib::GzipReader.new(StringIO.new(body.b)).read
         end
 
         def gzipped?(body)
@@ -418,10 +412,26 @@ def apply_headers(client, options)
           headers = options[:headers] || {}
           headers[CONTENT_TYPE_STR] = find_value(headers, CONTENT_TYPE_REGEX) || DEFAULT_CONTENT_TYPE
           headers[USER_AGENT_STR] = find_value(headers, USER_AGENT_REGEX) || user_agent_header(client)
-          client.headers[ACCEPT_ENCODING] = GZIP if use_compression?
+          client.headers[ACCEPT_ENCODING] = GZIP
+          client.headers[CONTENT_ENCODING] = GZIP if use_compression?
           client.headers.merge!(headers)
         end
 
+        def compress_request(body)
+          io = StringIO.new(''.b)
+          gzip_writer = Zlib::GzipWriter.new(io)
+          gzip_writer.write(body)
+          gzip_writer.close
+          io.string
+        end
+
+        def request_body(body)
+          return nil unless body
+
+          json_body = __convert_to_json(body)
+          use_compression? ? compress_request(json_body) : json_body
+        end
+
         def find_value(hash, regex)
           key_value = hash.find { |k, _v| k.to_s.downcase =~ regex }
           return unless key_value

lib/opensearch/transport/transport/http/curb.rb

@@ -49,7 +49,7 @@ def perform_request(method, path, params = {}, body = nil, headers = nil, opts =
                 connection.connection.set :nobody, true
               when 'GET', 'POST', 'PUT', 'DELETE'
                 connection.connection.set :nobody, false
-                connection.connection.put_data = __convert_to_json(body) if body
+                connection.connection.put_data = request_body(body) if body
                 if headers
                   if connection.connection.headers
                     connection.connection.headers.merge!(headers)

lib/opensearch/transport/transport/http/faraday.rb

@@ -58,7 +58,7 @@ def perform_request(method, path, params = {}, body = nil, headers = nil, opts =
               response = connection.connection.run_request(
                 method.downcase.to_sym,
                 url,
-                (body ? __convert_to_json(body) : nil),
+                request_body(body),
                 headers
               )
 

lib/opensearch/transport/transport/http/manticore.rb

@@ -91,7 +91,7 @@ def build_client(options = {})
           #
           def perform_request(method, path, params = {}, body = nil, headers = nil, opts = {})
             super do |connection, url|
-              params[:body] = __convert_to_json(body) if body
+              params[:body] = request_body(body) if body
               params[:headers] = headers if headers
               params = params.merge @request_options
               case method
@@ -167,7 +167,8 @@ def apply_headers(request_options, options)
             headers = (options && options[:headers]) || {}
             headers[CONTENT_TYPE_STR] = find_value(headers, CONTENT_TYPE_REGEX) || DEFAULT_CONTENT_TYPE
             headers[USER_AGENT_STR] = find_value(headers, USER_AGENT_REGEX) || user_agent_header
-            headers[ACCEPT_ENCODING] = GZIP if use_compression?
+            headers[ACCEPT_ENCODING] = GZIP
+            headers[CONTENT_ENCODING] = GZIP if use_compression?
             request_options.merge!(headers: headers)
           end
 

spec/opensearch/transport/base_spec.rb

@@ -310,4 +310,69 @@
       end
     end
   end
+
+  context 'when compression is enabled' do
+    let(:client) do
+      OpenSearch::Transport::Client.new(hosts: ['localhost'], compression: true)
+    end
+
+    let(:transport) { client.transport }
+    let(:connection) { transport.connections[0].connection }
+    let(:request_body) { { 'query' => { 'match_all' => {} } } }
+
+    before do
+      allow(connection).to receive(:run_request).and_return(
+        OpenStruct.new(status: 200, body: '{}', headers: { 'content-type' => 'application/json' })
+      )
+    end
+
+    it 'sets the Accept-Encoding and Content-Encoding headers' do
+      expect(connection.headers['Accept-Encoding']).to eq('gzip')
+      expect(connection.headers['Content-Encoding']).to eq('gzip')
+    end
+
+    it 'compresses the request body' do
+      expect(transport).to receive(:compress_request).and_call_original
+
+      transport.perform_request('POST', '_search', {}, request_body)
+    end
+
+    it 'decompresses the response body' do
+      expect(transport).to receive(:decompress_response).and_call_original
+
+      transport.perform_request('POST', '_search', {}, request_body)
+    end
+  end
+
+  context 'when compression is not enabled' do
+    let(:client) do
+      OpenSearch::Transport::Client.new(hosts: ['localhost'], compression: false)
+    end
+
+    let(:transport) { client.transport }
+    let(:connection) { transport.connections[0].connection }
+
+    before do
+      allow(connection).to receive(:run_request).and_return(
+        OpenStruct.new(status: 200, body: '{}', headers: { 'content-type' => 'application/json' })
+      )
+    end
+
+    it 'sets Accept-Encoding but not Content-Encoding' do
+      expect(connection.headers['Accept-Encoding']).to eq('gzip')
+      expect(connection.headers['Content-Encoding']).to be_nil
+    end
+
+    it 'does not compress the request body' do
+      expect(transport).not_to receive(:compress_request)
+
+      transport.perform_request('POST', '_search', {}, { 'query' => { 'match_all' => {} } })
+    end
+
+    it 'still decompresses the response body' do
+      expect(transport).to receive(:decompress_response).and_call_original
+
+      transport.perform_request('POST', '_search', {}, { 'query' => { 'match_all' => {} } })
+    end
+  end
 end

spec/opensearch/transport/client_spec.rb

@@ -1674,6 +1674,22 @@
           expect(client.perform_request('GET', '/').body).to be_a(Hash)
         end
       end
+
+      context 'when performing a search with a request body' do
+        it 'succeeds with compression enabled' do
+          client = described_class.new(hosts: OPENSEARCH_HOSTS, compression: true)
+
+          expect(client.transport.connections[0].connection.headers['Content-Encoding']).to eq('gzip')
+          expect(client.perform_request('POST', '_search', {}, { query: { match_all: {} } }).body).to be_a(Hash)
+        end
+
+        it 'succeeds with compression disabled' do
+          client = described_class.new(hosts: OPENSEARCH_HOSTS, compression: false)
+
+          expect(client.transport.connections[0].connection.headers['Content-Encoding']).to be_nil
+          expect(client.perform_request('POST', '_search', {}, { query: { match_all: {} } }).body).to be_a(Hash)
+        end
+      end
     end
 
     describe '#perform_request' do

spec/spec_helper.rb

@@ -37,6 +37,7 @@
 require 'jsonify'
 require 'yaml'
 require 'hashie/mash'
+require 'ostruct'
 
 require 'faraday/httpclient'
 require 'faraday/net_http_persistent'

test/transport/unit/transport_curb_test.rb

@@ -136,6 +136,64 @@ class OpenSearch::Transport::Transport::HTTP::FaradayTest < Minitest::Test
 
         assert_equal 'https://U:P@foobar:1234/', transport.connections.first.full_url('')
       end
+
+      context "when compression is enabled" do
+        setup do
+          @transport = Curb.new :hosts => [ { :host => 'foobar', :port => 1234 } ],
+                                :options => { :compression => true }
+        end
+
+        should "set Accept-Encoding and Content-Encoding headers" do
+          headers = @transport.connections.first.connection.headers
+          assert_equal 'gzip', headers['Accept-Encoding']
+          assert_equal 'gzip', headers['Content-Encoding']
+        end
+
+        should "compress the request body" do
+          @transport.expects(:compress_request).returns('compressed')
+          @transport.connections.first.connection.expects(:put_data=).with('compressed')
+          @transport.connections.first.connection.expects(:http).with(:POST).returns(stub_everything)
+          @transport.perform_request 'POST', '/', {}, { :foo => 'bar' }
+        end
+
+        should "decompress the response body" do
+          @transport.expects(:decompress_response).returns('{}')
+          @transport.connections.first.connection.expects(:http).with(:POST).returns(stub_everything)
+          @transport.perform_request 'POST', '/', {}, { :foo => 'bar' }
+        end
+      end
+
+      context "when compression is not enabled" do
+        setup do
+          @transport = Curb.new :hosts => [ { :host => 'foobar', :port => 1234 } ],
+                                :options => { :compression => false }
+        end
+
+        should "set Accept-Encoding but not Content-Encoding" do
+          headers = @transport.connections.first.connection.headers
+          assert_equal 'gzip', headers['Accept-Encoding']
+          assert_nil headers['Content-Encoding']
+        end
+
+        should "not compress the request body" do
+          @transport.expects(:compress_request).never
+          @transport.connections.first.connection.expects(:http).with(:POST).returns(stub_everything)
+          @transport.perform_request 'POST', '/', {}, { :foo => 'bar' }
+        end
+
+        should "still decompress the response body" do
+          @transport.expects(:decompress_response).returns('{}')
+          @transport.connections.first.connection.expects(:http).with(:POST).returns(stub_everything)
+          @transport.perform_request 'POST', '/', {}, { :foo => 'bar' }
+        end
+
+        should "not call put_data= without a request body" do
+          # Curl::Easy#put_data= enables upload mode even for nil - it should not be invoked even for an empty body
+          @transport.connections.first.connection.expects(:put_data=).never
+          @transport.connections.first.connection.expects(:http).with(:GET).returns(stub_everything)
+          @transport.perform_request 'GET', '/'
+        end
+      end
     end
 
   end

test/transport/unit/transport_faraday_test.rb

@@ -232,6 +232,56 @@ class OpenSearch::Transport::Transport::HTTP::FaradayTest < Minitest::Test
 
       assert_equal 'https://U:P@foobar:1234/', transport.connections.first.full_url('')
     end
+
+    context "when compression is enabled" do
+      setup do
+        @transport = Faraday.new :hosts => [ { :host => 'foobar', :port => 1234 } ],
+                                 :options => { :compression => true }
+      end
+
+      should "set Accept-Encoding and Content-Encoding headers" do
+        headers = @transport.connections.first.connection.headers
+        assert_equal 'gzip', headers['Accept-Encoding']
+        assert_equal 'gzip', headers['Content-Encoding']
+      end
+
+      should "compress the request body" do
+        @transport.expects(:compress_request).returns('compressed')
+        @transport.connections.first.connection.expects(:run_request).returns(stub_everything)
+        @transport.perform_request 'POST', '/', {}, { :foo => 'bar' }
+      end
+
+      should "decompress the response body" do
+        @transport.expects(:decompress_response).returns('{}')
+        @transport.connections.first.connection.expects(:run_request).returns(stub_everything)
+        @transport.perform_request 'POST', '/', {}, { :foo => 'bar' }
+      end
+    end
+
+    context "when compression is not enabled" do
+      setup do
+        @transport = Faraday.new :hosts => [ { :host => 'foobar', :port => 1234 } ],
+                                 :options => { :compression => false }
+      end
+
+      should "set Accept-Encoding but not Content-Encoding" do
+        headers = @transport.connections.first.connection.headers
+        assert_equal 'gzip', headers['Accept-Encoding']
+        assert_nil headers['Content-Encoding']
+      end
+
+      should "not compress the request body" do
+        @transport.expects(:compress_request).never
+        @transport.connections.first.connection.expects(:run_request).returns(stub_everything)
+        @transport.perform_request 'POST', '/', {}, { :foo => 'bar' }
+      end
+
+      should "still decompress the response body" do
+        @transport.expects(:decompress_response).returns('{}')
+        @transport.connections.first.connection.expects(:run_request).returns(stub_everything)
+        @transport.perform_request 'POST', '/', {}, { :foo => 'bar' }
+      end
+    end
   end
 
 end