Add Stream option to PerformRequestOptions (#1478)

Add a `Stream` option to `PerformRequestOptions` which allows
direct access to the HTTP response body.

Close #1444

Author: atombender

client.go

@@ -1286,6 +1286,7 @@ type PerformRequestOptions struct {
 	RetryStatusCodes []int
 	Headers          http.Header
 	MaxResponseSize  int64
+	Stream           bool
 }
 
 // PerformRequest does a HTTP request to Elasticsearch.
@@ -1294,6 +1295,9 @@ type PerformRequestOptions struct {
 // Optionally, a list of HTTP error codes to ignore can be passed.
 // This is necessary for services that expect e.g. HTTP status 404 as a
 // valid outcome (Exists, IndicesExists, IndicesTypeExists).
+//
+// If Stream is set, the returned BodyReader field must be closed, even
+// if PerformRequest returns an error.
 func (c *Client) PerformRequest(ctx context.Context, opt PerformRequestOptions) (*Response, error) {
 	start := time.Now().UTC()
 
@@ -1448,7 +1452,10 @@ func (c *Client) PerformRequest(ctx context.Context, opt PerformRequestOptions)
 				continue // try again
 			}
 		}
-		defer res.Body.Close()
+
+		if !opt.Stream {
+			defer res.Body.Close()
+		}
 
 		// Tracing
 		c.dumpResponse(res)
@@ -1465,14 +1472,14 @@ func (c *Client) PerformRequest(ctx context.Context, opt PerformRequestOptions)
 		if err := checkResponse((*http.Request)(req), res, opt.IgnoreErrors...); err != nil {
 			// No retry if request succeeded
 			// We still try to return a response.
-			resp, _ = c.newResponse(res, opt.MaxResponseSize)
+			resp, _ = c.newResponse(res, opt.MaxResponseSize, opt.Stream)
 			return resp, err
 		}
 
 		// We successfully made a request with this connection
 		conn.MarkAsHealthy()
 
-		resp, err = c.newResponse(res, opt.MaxResponseSize)
+		resp, err = c.newResponse(res, opt.MaxResponseSize, opt.Stream)
 		if err != nil {
 			return nil, err
 		}

client_test.go

@@ -859,7 +859,7 @@ func TestClientSniffTimeoutLeak(t *testing.T) {
 	cli := &Client{
 		c: &http.Client{},
 		conns: []*conn{
-			&conn{
+			{
 				url: "http://" + addr + "/",
 			},
 		},
@@ -1161,6 +1161,38 @@ func TestPerformRequest(t *testing.T) {
 	}
 }
 
+func TestPerformRequestWithStream(t *testing.T) {
+	client, err := NewClient()
+	if err != nil {
+		t.Fatal(err)
+	}
+	res, err := client.PerformRequest(context.TODO(), PerformRequestOptions{
+		Method: "GET",
+		Path:   "/",
+		Stream: true,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if res == nil {
+		t.Fatal("expected response to be != nil")
+	}
+	if res.BodyReader == nil {
+		t.Fatal("expected res.BodyReader to be != nil")
+	}
+	if res.Body != nil {
+		t.Fatal("expected res.Body to be == nil")
+	}
+
+	ret := new(PingResult)
+	if err := json.NewDecoder(res.BodyReader).Decode(ret); err != nil {
+		t.Fatalf("expected no error on decode; got: %v", err)
+	}
+	if ret.ClusterName == "" {
+		t.Errorf("expected cluster name; got: %q", ret.ClusterName)
+	}
+}
+
 func TestPerformRequestWithSimpleClient(t *testing.T) {
 	client, err := NewSimpleClient()
 	if err != nil {

response.go

@@ -24,21 +24,25 @@ type Response struct {
 	// Header is the HTTP header from the HTTP response.
 	// Keys in the map are canonicalized (see http.CanonicalHeaderKey).
 	Header http.Header
-	// Body is the deserialized response body.
+	// Body is the deserialized response body. Only available if streaming is disabled.
 	Body json.RawMessage
 	// DeprecationWarnings lists all deprecation warnings returned from
 	// Elasticsearch.
 	DeprecationWarnings []string
+	// BodyReader is the body as a reader. Only available if streaming is enabled.
+	BodyReader io.ReadCloser
 }
 
 // newResponse creates a new response from the HTTP response.
-func (c *Client) newResponse(res *http.Response, maxBodySize int64) (*Response, error) {
+func (c *Client) newResponse(res *http.Response, maxBodySize int64, stream bool) (*Response, error) {
 	r := &Response{
 		StatusCode:          res.StatusCode,
 		Header:              res.Header,
 		DeprecationWarnings: res.Header["Warning"],
 	}
-	if res.Body != nil {
+	if stream {
+		r.BodyReader = res.Body
+	} else if res.Body != nil {
 		body := io.Reader(res.Body)
 		if maxBodySize > 0 {
 			if res.ContentLength > maxBodySize {

response_test.go

@@ -29,7 +29,7 @@ func BenchmarkResponse(b *testing.B) {
 			StatusCode: http.StatusOK,
 		}
 		var err error
-		resp, err = c.newResponse(res, 0)
+		resp, err = c.newResponse(res, 0, false)
 		if err != nil {
 			b.Fatal(err)
 		}