Add BulkIndexerItem.RequireDataStream (#235)

If set to true, require_data_stream=true
will be added to the bulk request item action.

Author: axw

bulk_indexer.go

@@ -286,12 +286,13 @@ func (b *BulkIndexer) BytesUncompressedFlushed() int {
 }
 
 type BulkIndexerItem struct {
-	Index            string
-	DocumentID       string
-	Pipeline         string
-	Action           string
-	Body             io.WriterTo
-	DynamicTemplates map[string]string
+	Index             string
+	DocumentID        string
+	Pipeline          string
+	Action            string
+	Body              io.WriterTo
+	DynamicTemplates  map[string]string
+	RequireDataStream bool
 }
 
 // Add encodes an item in the buffer.
@@ -307,7 +308,14 @@ func (b *BulkIndexer) Add(item BulkIndexerItem) error {
 		return fmt.Errorf("%s is not a valid action", action)
 	}
 
-	b.writeMeta(item.Index, item.DocumentID, item.Pipeline, action, item.DynamicTemplates)
+	b.writeMeta(
+		item.Index,
+		item.DocumentID,
+		item.Pipeline,
+		action,
+		item.DynamicTemplates,
+		item.RequireDataStream,
+	)
 	if _, err := item.Body.WriteTo(b.writer); err != nil {
 		return fmt.Errorf("failed to write bulk indexer item: %w", err)
 	}
@@ -318,7 +326,11 @@ func (b *BulkIndexer) Add(item BulkIndexerItem) error {
 	return nil
 }
 
-func (b *BulkIndexer) writeMeta(index, documentID, pipeline, action string, dynamicTemplates map[string]string) {
+func (b *BulkIndexer) writeMeta(
+	index, documentID, pipeline, action string,
+	dynamicTemplates map[string]string,
+	requireDataStream bool,
+) {
 	b.jsonw.RawString(`{"`)
 	b.jsonw.RawString(action)
 	b.jsonw.RawString(`":{`)
@@ -363,6 +375,13 @@ func (b *BulkIndexer) writeMeta(index, documentID, pipeline, action string, dyna
 		b.jsonw.RawByte('}')
 		first = false
 	}
+	if requireDataStream {
+		if !first {
+			b.jsonw.RawByte(',')
+		}
+		b.jsonw.RawString(`"require_data_stream":true`)
+		first = false
+	}
 	b.jsonw.RawString("}}\n")
 	b.writer.Write(b.jsonw.Bytes())
 	b.jsonw.Reset()

bulk_indexer_test.go

@@ -21,6 +21,8 @@ import (
 	"context"
 	"encoding/json"
 	"net/http"
+	"strconv"
+	"strings"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -266,3 +268,30 @@ func TestAction(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, int64(3), stat.Indexed)
 }
+
+func TestItemRequireDataStream(t *testing.T) {
+	client := docappendertest.NewMockElasticsearchClient(t, func(w http.ResponseWriter, r *http.Request) {
+		_, meta, result, _ := docappendertest.DecodeBulkRequestWithStatsAndMeta(r)
+		require.Len(t, meta, 2)
+		assert.False(t, meta[0].RequireDataStream)
+		assert.True(t, meta[1].RequireDataStream)
+		json.NewEncoder(w).Encode(result)
+	})
+	indexer, err := docappender.NewBulkIndexer(docappender.BulkIndexerConfig{
+		Client: client,
+	})
+	require.NoError(t, err)
+
+	for _, required := range []bool{false, true} {
+		err := indexer.Add(docappender.BulkIndexerItem{
+			Index:             strconv.FormatBool(required),
+			Body:              strings.NewReader(`{}`),
+			RequireDataStream: required,
+		})
+		require.NoError(t, err)
+	}
+
+	stat, err := indexer.Flush(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, int64(2), stat.Indexed)
+}

docappendertest/docappendertest.go

@@ -39,6 +39,15 @@ import (
 	"github.com/elastic/go-elasticsearch/v8/esutil"
 )
 
+type BulkRequestItemMeta struct {
+	Action            string            `json:"-"`
+	Index             string            `json:"_index"`
+	DocumentID        string            `json:"_id"`
+	Pipeline          string            `json:"pipeline"`
+	DynamicTemplates  map[string]string `json:"dynamic_templates"`
+	RequireDataStream bool              `json:"require_data_stream"`
+}
+
 // TimestampFormat holds the time format for formatting timestamps according to
 // Elasticsearch's strict_date_optional_time date format, which includes a fractional
 // seconds component.
@@ -60,6 +69,18 @@ func DecodeBulkRequestWithStats(r *http.Request) (
 	return indexed, result, stats
 }
 
+// DecodeBulkRequestWithStatsAndMeta decodes a /_bulk request's body,
+// returning the decoded bulk request action/meta and documents,
+// and a response body and stats about the request.
+func DecodeBulkRequestWithStatsAndMeta(r *http.Request) (
+	docs [][]byte,
+	meta []BulkRequestItemMeta,
+	res esutil.BulkIndexerResponse,
+	stats RequestStats,
+) {
+	return decodeBulkRequest(r)
+}
+
 // DecodeBulkRequestWithStatsAndDynamicTemplates decodes a /_bulk request's body,
 // returning the decoded documents and a response body and stats about request, and per-request dynamic templates.
 func DecodeBulkRequestWithStatsAndDynamicTemplates(r *http.Request) (
@@ -79,7 +100,22 @@ func DecodeBulkRequestWithStatsAndDynamicTemplatesAndPipelines(r *http.Request)
 	res esutil.BulkIndexerResponse,
 	stats RequestStats,
 	dynamicTemplates []map[string]string,
-	pipelines []string) {
+	pipelines []string,
+) {
+	docs, meta, res, stats := decodeBulkRequest(r)
+	for _, meta := range meta {
+		dynamicTemplates = append(dynamicTemplates, meta.DynamicTemplates)
+		pipelines = append(pipelines, meta.Pipeline)
+	}
+	return docs, res, stats, dynamicTemplates, pipelines
+}
+
+func decodeBulkRequest(r *http.Request) (
+	docs [][]byte,
+	meta []BulkRequestItemMeta,
+	result esutil.BulkIndexerResponse,
+	stats RequestStats,
+) {
 	body := r.Body
 	switch r.Header.Get("Content-Encoding") {
 	case "gzip":
@@ -95,15 +131,10 @@ func DecodeBulkRequestWithStatsAndDynamicTemplatesAndPipelines(r *http.Request)
 	}
 	body = cr
 	defer cr.Close()
+
 	scanner := bufio.NewScanner(body)
-	var indexed [][]byte
-	var result esutil.BulkIndexerResponse
 	for scanner.Scan() {
-		action := make(map[string]struct {
-			Index            string            `json:"_index"`
-			DynamicTemplates map[string]string `json:"dynamic_templates"`
-			Pipeline         string            `json:"pipeline"`
-		})
+		action := make(map[string]BulkRequestItemMeta)
 		if err := json.NewDecoder(strings.NewReader(scanner.Text())).Decode(&action); err != nil {
 			panic(err)
 		}
@@ -118,14 +149,16 @@ func DecodeBulkRequestWithStatsAndDynamicTemplatesAndPipelines(r *http.Request)
 		if !json.Valid(doc) {
 			panic(fmt.Errorf("invalid JSON: %s", doc))
 		}
-		indexed = append(indexed, doc)
+		docs = append(docs, doc)
 
 		item := esutil.BulkIndexerResponseItem{Status: http.StatusCreated, Index: action[actionType].Index}
 		result.Items = append(result.Items, map[string]esutil.BulkIndexerResponseItem{actionType: item})
-		dynamicTemplates = append(dynamicTemplates, action[actionType].DynamicTemplates)
-		pipelines = append(pipelines, action[actionType].Pipeline)
+
+		itemMeta := action[actionType]
+		itemMeta.Action = actionType
+		meta = append(meta, itemMeta)
 	}
-	return indexed, result, RequestStats{int64(cr.bytesRead)}, dynamicTemplates, pipelines
+	return docs, meta, result, RequestStats{int64(cr.bytesRead)}
 }
 
 // NewMockElasticsearchClient returns an elasticsearch.Client which sends /_bulk requests to bulkHandler.