fix(client): re-encode _search body on 401 retry (empty-body / dropped-payload bug) (#118)

* fix(client): re-encode _search body on 401 retry and check transport error

On an HTTP 401 the connector re-authenticates and replays the SAME
*esapi.SearchRequest via req.Do. The request body is an io.Reader that the
first req.Do fully drains, so the retried _search was sent with an empty body.
Elasticsearch interprets an empty _search body as match_all and returns
unfiltered results — the intermittent "dropped payload on retry" symptom seen
after token expiry/rotation mid-traffic (JPMorgan Chase support ticket #15000).

- Capture the encoded body once and rebuild a fresh reader before every
  attempt so the retried _search carries the identical query body.
- Check the error from the first req.Do BEFORE calling res.IsError(), which
  would panic on a nil response after a transport-level failure.
- Add per-attempt request logging of the actual body + target index sent on
  the wire; the retry line carries the literal "Retry Query" marker. The log
  reflects what is truly sent, so it shows the empty body under the old bug.
- Add hermetic functional tests (httptest fake ES, no real cluster) that
  reproduce the empty-retry-body bug and verify the fix. The env var
  ELASTICSEARCH_DISABLE_RETRY_BODY_REBUILD flips between buggy and fixed paths.

Co-authored-by: Mohd Bilal <bilal@hasura.io>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* docs: add pr_result.md summarizing the 401-retry fix, tests and logging

Co-authored-by: Mohd Bilal <bilal@hasura.io>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* refactor(client): remove testing-only env-var toggle for retry body rebuild

Address review feedback: drop the ELASTICSEARCH_DISABLE_RETRY_BODY_REBUILD
mechanism entirely. The retry now always rebuilds the request body
unconditionally before req.Do, and the "Retry Query" log line logs the exact
bytes sent (string(body)).

- Delete retryBodyRebuildDisabled() and the guard around the rebuild.
- Delete retryBodyForLog(); log string(body) which equals the wire body.
- Tests no longer reference the env var; the functional test asserts only the
  fixed behaviour. Reproducing the old bug is documented as commenting out the
  rebuild line in search().

Primary fix (rebuild body before every attempt), secondary fix (check the first
req.Do error before res.IsError()), and per-attempt request logging with the
"Retry Query" marker are all retained.

Co-authored-by: Mohd Bilal <bilal@hasura.io>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* chore: remove stray pr_result.md from repo

This reporting file was committed by mistake and should not live in the repo.

Co-authored-by: Mohd Bilal <bilal@hasura.io>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: dliub <noreply@anthropic.com>
Co-authored-by: Mohd Bilal <bilal@hasura.io>

Author: 3 people

elasticsearch/client.go

@@ -6,7 +6,9 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"os"
+	"strings"
 
 	"github.com/elastic/go-elasticsearch/v8"
 	"github.com/elastic/go-elasticsearch/v8/esapi"
@@ -166,21 +168,48 @@ func (e *Client) Search(ctx context.Context, index string, body map[string]inter
 
 // search is a helper function to perform a search operation in elastic search.
 func (e *Client) search(ctx context.Context, o ...func(*esapi.SearchRequest)) (*esapi.Response, error) {
+	logger := connector.GetLogger(ctx)
+
 	req := &esapi.SearchRequest{}
 
 	for _, opt := range o {
 		opt(req)
 	}
 
+	// req.Body is an io.Reader that is fully drained by req.Do. If we retry the
+	// request after a 401 (see below) by calling req.Do again on the same req,
+	// the already-drained reader sends an empty body. Elasticsearch treats an
+	// empty _search body as a match_all query and returns unfiltered results.
+	// To make the request safely repeatable, capture the encoded body once and
+	// rebuild a fresh reader before every attempt.
+	body, err := drainBody(req.Body)
+	if err != nil {
+		return nil, fmt.Errorf("error reading search request body: %w", err)
+	}
+	index := strings.Join(req.Index, ",")
+
+	// First attempt.
+	req.Body = bytes.NewReader(body)
+	logger.DebugContext(ctx, "Query", "index", index, "body", string(body))
 	res, err := req.Do(ctx, e.client)
+	// Check the transport error before touching res: on a transport-level
+	// failure res is nil and res.IsError() would panic.
+	if err != nil {
+		return nil, fmt.Errorf("error while querying: %s", err)
+	}
 
 	if res.IsError() {
 		if res.StatusCode == 401 {
-			// Unauthorized error, reauthenticate and retry
-			err = e.Reauthenticate(ctx)
-			if err != nil {
+			// Unauthorized error, reauthenticate and retry.
+			if err = e.Reauthenticate(ctx); err != nil {
 				return nil, fmt.Errorf("error: %s", err)
 			}
+			// Rebuild the body so the retried request carries the same query
+			// instead of the already-drained (empty) reader. The retried reader
+			// holds exactly these bytes, so logging string(body) reflects what is
+			// actually sent over the wire.
+			req.Body = bytes.NewReader(body)
+			logger.DebugContext(ctx, "Retry Query", "index", index, "body", string(body))
 			res, err = req.Do(ctx, e.client)
 			if err != nil {
 				return nil, fmt.Errorf("error: %s", err)
@@ -192,6 +221,15 @@ func (e *Client) search(ctx context.Context, o ...func(*esapi.SearchRequest)) (*
 	return res, err
 }
 
+// drainBody reads an io.Reader fully and returns its bytes, handling a nil
+// reader (no body) gracefully.
+func drainBody(r io.Reader) ([]byte, error) {
+	if r == nil {
+		return nil, nil
+	}
+	return io.ReadAll(r)
+}
+
 // Explain performs a search with explain operation in elastic search.
 //
 // Since the Explain API requires document ID, we can't use it.

elasticsearch/client_retry_test.go

@@ -0,0 +1,238 @@
+package elasticsearch
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+)
+
+// newFakeES returns an httptest server that emulates an Elasticsearch cluster
+// which fails the first _search with 401 (token expired mid-traffic) and then
+// succeeds. It records the body received on every _search attempt into bodies.
+func newFakeES(t *testing.T, bodies *[]string, mu *sync.Mutex) *httptest.Server {
+	t.Helper()
+	var attempts int
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("X-Elastic-Product", "Elasticsearch")
+		switch {
+		case r.Method == http.MethodHead && r.URL.Path == "/":
+			w.WriteHeader(http.StatusOK)
+		case r.URL.Path == "/" && r.Method == http.MethodGet:
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write([]byte(`{"version":{"number":"8.0.0"}}`))
+		default:
+			bodyBytes, _ := io.ReadAll(r.Body)
+			mu.Lock()
+			attempts++
+			attempt := attempts
+			*bodies = append(*bodies, string(bodyBytes))
+			mu.Unlock()
+			if attempt == 1 {
+				w.WriteHeader(http.StatusUnauthorized)
+				_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
+				return
+			}
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write([]byte(`{"took":1,"hits":{"total":{"value":0},"hits":[]}}`))
+		}
+	}))
+}
+
+// TestSearchRetryBodyOn401 is a hermetic functional test for the "dropped
+// payload on intermittent retry" bug (JPMorgan Chase support ticket #15000).
+//
+// It stands up a fake Elasticsearch server that returns 401 on the first
+// _search call (simulating a token expiring mid-traffic) and 200 afterwards.
+// It captures the request body received on EVERY _search attempt and asserts
+// that the retry, sent after re-authentication, carries the SAME, non-empty
+// query body as the first attempt.
+//
+//	go test -v -run TestSearchRetryBodyOn401 ./elasticsearch/
+//
+// To manually reproduce the OLD buggy behaviour (retry sends an empty body),
+// comment out the body-rebuild line in search() in client.go:
+//
+//	// req.Body = bytes.NewReader(body)   // <- comment this out before the retry req.Do
+//
+// and re-run this test: it will fail with "BUG REPRODUCED", because the retry
+// then reuses the already-drained reader and sends an empty body.
+func TestSearchRetryBodyOn401(t *testing.T) {
+	var (
+		mu           sync.Mutex
+		searchBodies []string
+	)
+
+	server := newFakeES(t, &searchBodies, &mu)
+	defer server.Close()
+
+	// Point the connector at the fake server. These are required by
+	// getConfigFromEnv so that Reauthenticate succeeds against the fake server.
+	t.Setenv("ELASTICSEARCH_URL", server.URL)
+	t.Setenv("ELASTICSEARCH_USERNAME", "elastic")
+	t.Setenv("ELASTICSEARCH_PASSWORD", "changeme")
+
+	ctx := context.Background()
+	client, err := NewClient(ctx)
+	if err != nil {
+		t.Fatalf("failed to create client: %v", err)
+	}
+
+	query := map[string]interface{}{
+		"query": map[string]interface{}{
+			"term": map[string]interface{}{"customer_id": "JPMC-42"},
+		},
+		"size": 10,
+	}
+	expectedBody, err := json.Marshal(query)
+	if err != nil {
+		t.Fatalf("failed to marshal query: %v", err)
+	}
+
+	if _, err := client.Search(ctx, "transactions", query); err != nil {
+		t.Fatalf("Search returned error: %v", err)
+	}
+
+	mu.Lock()
+	defer mu.Unlock()
+
+	if len(searchBodies) != 2 {
+		t.Fatalf("expected exactly 2 _search attempts (initial + retry), got %d", len(searchBodies))
+	}
+
+	// The first attempt must always carry the full query body.
+	if !jsonEqual(t, searchBodies[0], string(expectedBody)) {
+		t.Fatalf("first attempt body mismatch:\n  got:  %q\n  want: %q", searchBodies[0], expectedBody)
+	}
+
+	// The crux of the bug: the retry must carry the SAME, non-empty body.
+	retryBody := searchBodies[1]
+	if len(trimWS(retryBody)) == 0 {
+		t.Fatalf("BUG REPRODUCED: retry sent an EMPTY body (Elasticsearch would treat this as match_all and return unfiltered results). retry body=%q", retryBody)
+	}
+	if !jsonEqual(t, retryBody, string(expectedBody)) {
+		t.Fatalf("retry body mismatch:\n  got:  %q\n  want: %q", retryBody, expectedBody)
+	}
+
+	t.Logf("OK: both attempts sent identical, non-empty query body: %s", expectedBody)
+}
+
+// TestSearchPerAttemptLogging verifies the per-attempt request logging: the
+// connector logs the actual _search body and target index it sends on EVERY
+// attempt, and the retry log line carries the literal "Retry Query" marker so
+// retries are easy to grep. The logged body reflects what is ACTUALLY sent, so
+// the "Retry Query" line carries the same non-empty body as the first attempt.
+//
+//	go test -v -run TestSearchPerAttemptLogging ./elasticsearch/
+func TestSearchPerAttemptLogging(t *testing.T) {
+	var (
+		mu           sync.Mutex
+		searchBodies []string
+	)
+	server := newFakeES(t, &searchBodies, &mu)
+	defer server.Close()
+
+	t.Setenv("ELASTICSEARCH_URL", server.URL)
+	t.Setenv("ELASTICSEARCH_USERNAME", "elastic")
+	t.Setenv("ELASTICSEARCH_PASSWORD", "changeme")
+
+	// Capture logs: GetLogger falls back to slog.Default() when no logger is set
+	// on the context, so redirect the default logger to a debug-level buffer.
+	var logBuf bytes.Buffer
+	prev := slog.Default()
+	slog.SetDefault(slog.New(slog.NewJSONHandler(&logBuf, &slog.HandlerOptions{Level: slog.LevelDebug})))
+	defer slog.SetDefault(prev)
+
+	ctx := context.Background()
+	client, err := NewClient(ctx)
+	if err != nil {
+		t.Fatalf("failed to create client: %v", err)
+	}
+
+	query := map[string]interface{}{
+		"query": map[string]interface{}{"term": map[string]interface{}{"customer_id": "JPMC-42"}},
+	}
+	if _, err := client.Search(ctx, "transactions", query); err != nil {
+		t.Fatalf("Search returned error: %v", err)
+	}
+
+	logs := logBuf.String()
+	t.Logf("captured logs:\n%s", logs)
+
+	// Per-attempt logging: first attempt logged as "Query".
+	if !strings.Contains(logs, `"msg":"Query"`) {
+		t.Errorf(`expected a per-attempt "Query" log line, not found in:\n%s`, logs)
+	}
+	// Retry attempt logged with the literal "Retry Query" marker.
+	if !strings.Contains(logs, `"msg":"Retry Query"`) {
+		t.Errorf(`expected a "Retry Query" marked log line, not found in:\n%s`, logs)
+	}
+	// The log must carry the target index.
+	if !strings.Contains(logs, `"index":"transactions"`) {
+		t.Errorf(`expected the target index in the log line, not found in:\n%s`, logs)
+	}
+
+	// The logged retry body must reflect what is ACTUALLY sent on the wire: a
+	// non-empty body identical to the first attempt.
+	retryBody := logFieldForMsg(t, logs, "Retry Query", "body")
+	if trimWS(retryBody) == "" {
+		t.Fatalf("expected a non-empty logged retry body, got %q", retryBody)
+	}
+	if !jsonEqual(t, retryBody, logFieldForMsg(t, logs, "Query", "body")) {
+		t.Errorf("logged retry body %q does not match first attempt body", retryBody)
+	}
+	t.Logf("retry logged the full body: %s", retryBody)
+}
+
+// logFieldForMsg returns the value of field for the first JSON log record whose
+// "msg" equals msg.
+func logFieldForMsg(t *testing.T, logs, msg, field string) string {
+	t.Helper()
+	for _, line := range strings.Split(strings.TrimSpace(logs), "\n") {
+		if line == "" {
+			continue
+		}
+		var rec map[string]interface{}
+		if err := json.Unmarshal([]byte(line), &rec); err != nil {
+			continue
+		}
+		if rec["msg"] == msg {
+			if v, ok := rec[field].(string); ok {
+				return v
+			}
+		}
+	}
+	return ""
+}
+
+// jsonEqual compares two JSON documents for semantic equality (ignoring key
+// order and insignificant whitespace).
+func jsonEqual(t *testing.T, a, b string) bool {
+	t.Helper()
+	var av, bv interface{}
+	if err := json.Unmarshal([]byte(a), &av); err != nil {
+		return false
+	}
+	if err := json.Unmarshal([]byte(b), &bv); err != nil {
+		return false
+	}
+	ab, _ := json.Marshal(av)
+	bb, _ := json.Marshal(bv)
+	return string(ab) == string(bb)
+}
+
+func trimWS(s string) string {
+	out := make([]rune, 0, len(s))
+	for _, r := range s {
+		if r != ' ' && r != '\n' && r != '\t' && r != '\r' {
+			out = append(out, r)
+		}
+	}
+	return string(out)
+}