Repository: olivere/elastic (go)

Issue:
Direct access to HTTP response
Background: We're using the client in an application where concurrent memory usage is a significant concern, and queries can return many megabytes of data in some cases. To reduce memory usage and increase pipelining of throughout our query handling, we've developed a solution where we stream the byte stream from Elasticsearch instead of unmarshalling of the entire response into a struct.

This works really well, but we had to use a custom JSON parser to accomplish this, and we also needed direct access to the HTTP response instead of getting a struct back. Unfortunately, the client doesn't permit this; we encountered some roadblocks:

* Every `Do()` method unmarshals the response.
* The client has a `PerformRequest()`, but it reads the entire body and does not return a stream.
* There's no way to ask a client for its configuration so that one can set up a new request
* All the logic to build an HTTP request are buried in the `Do()` methods

We ended up creating a "wrapper client" that contains an `*elastic.Client`, but also carries around with it the URL of the server and some other configuration. Then we use `net/http` to make an HTTP request, bypassing the client altogether. Fortunately, we are able to reuse structs like `elastic.SearchRequest`, but we have to construct the HTTP envelope (e.g. `scroll`) ourselves. We also have to do things like URL construction, logging, status handling, retrying, etc. ourselves.

Here's a snippet from our current logic, greatly simplified:

```go
b, err := search.Body()  // search is a *elastic.SearchRequest we've prepared
if err != nil {
	return err
}

url := baseURL
url.Query = url.Values{"scroll": []string{keepAlive}}

req, err := http.NewRequest(http.MethodPost, url.String(), bytes.NewReader(b))
if err != nil {
	...
}

resp, err := http.DefaultClient.Do(ctx, req)
if err != nil {
	...
}

return newJSONStream(resp.Body)
```

It would be much better if we could do something like:

```go
scroll := client.Scroll().Index(indeXName).KeepAlive(keepAlive) // etc.

resp, err := scroll.DoStream(ctx)  // Returns an *http.Response
if err != nil {
	...
}

return newJSONStream(resp.Body)
```

Or at least something like:

```
resp, err := client.PerformRequest(ctx, elastic.PerformRequestOptions{
	Body: search,
	Params: url.Values{"scroll": []string{keepAlive}},
	Stream: true,  // Signals that body should be returned as a stream
})
if err != nil {
	...
}

return newJSONStream(resp.Body)
```

Thoughts?

Pull request:
Add stream option to PerformRequest()
Add an option `Stream` to `PerformRequestOptions`, allowing a response to be parsed as a stream. Fixes #1444.

Patch:
diff --git a/client.go b/client.go
index bd2f2fe..8752302 100644
--- a/client.go
+++ b/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
 		}
diff --git a/response.go b/response.go
index 200ed3d..3e0c7de 100644
--- a/response.go
+++ b/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 {


Test patch:
diff --git a/client_test.go b/client_test.go
index a87e9b7..883ae29 100644
--- a/client_test.go
+++ b/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 {
diff --git a/response_test.go b/response_test.go
index 74552b7..3e3a6e6 100644
--- a/response_test.go
+++ b/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)
 		}