excon_backend.rb in riak-client-0.9.2

- old
+ new

@@ -9,11 +9,14 @@
 #    Unless required by applicable law or agreed to in writing, software
 #    distributed under the License is distributed on an "AS IS" BASIS,
 #    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
-require 'riak'
+
+require 'riak/failed_request'
+require 'riak/client/http_backend'
+require 'riak/client/http_backend/request_headers'
 require 'riak/client/pump'
 
 module Riak
   class Client
     # An HTTP backend for Riak::Client that uses Wesley Beary's Excon
@@ -40,21 +43,26 @@
         }
         params[:query] = uri.query if uri.query
         params[:body] = data if [:put,:post].include?(method)
         params[:idempotent] = (method != :post)
 
-        block = Pump.new(block) if block_given?
+        if block_given?
+          pump = Pump.new(block)
+          # Later versions of Excon pass multiple arguments to the block
+          block = lambda {|*args| pump.pump(args.first) }
+        end
 
         response = connection.request(params, &block)
+        response_headers.initialize_http_header(response.headers)
+
         if valid_response?(expect, response.status)
-          response_headers.initialize_http_header(response.headers)
           result = {:headers => response_headers.to_hash, :code => response.status}
           if return_body?(method, response.status, block_given?)
             result[:body] = response.body
           end
           result
         else
-          raise HTTPFailedRequest.new(method, expect, response.status, response.headers, response.body)
+          raise HTTPFailedRequest.new(method, expect, response.status, response_headers.to_hash, response.body)
         end
       end
 
       def connection
         @connection ||= Excon::Connection.new(root_uri.to_s)