Repository: elastic/elasticsearch-js (ts)

Issue:
Bulk helper document manipulation
Currently, you can't manipulate the document that gets passed to the bulk helper, if you need to do that, you should pipe the datasource in a transform stream or inside an async generator before giving it to the bulk helper:
```js
const { createReadStream } = require('fs')
const split = require('split2')
const { Client } = require('@elastic/elasticsearch')

async function * manipulate (stream) {
  for await (const chunk of stream) {
    // chunk manipulation
    yield chunk
  }
}

const stream = createReadStream('./dataset.ndjson').pipe(split())

const client = new Client({ node: 'http://localhost:9200' })
const result = await client.helpers.bulk({
  datasource: manipulate(stream),
  onDocument (doc) {
    return {
      index: { _index: 'my-index' }
    }
  }
})
```

While the solution above works perfectly, it does not support the case where the document contains metadata that does not need to be indexed, such as the index name for example:
```js
const { createReadStream } = require('fs')
const split = require('split2')
const { Client } = require('@elastic/elasticsearch')

// assuming the document has the following shape:
// { index: string, id: string, body: object }

const client = new Client({ node: 'http://localhost:9200' })
const result = await client.helpers.bulk({
  datasource: createReadStream('./dataset.ndjson').pipe(split(JSON.parse)),
  onDocument (doc) {
    return {
      index: { _index: doc.index, _id: doc.id }
    }
  }
})
```
In the code above we are indexing the document metadata as well.
For solving this problem, we could introduce another callback, to allow the document manipulation after the creation of the bulk action:
```js
const { createReadStream } = require('fs')
const split = require('split2')
const { Client } = require('@elastic/elasticsearch')

// assuming the document has the following shape:
// { index: string, id: string, body: object }

const client = new Client({ node: 'http://localhost:9200' })
const result = await client.helpers.bulk({
  datasource: createReadStream('./dataset.ndjson').pipe(split(JSON.parse)),
  onDocument (doc) {
    return {
      index: { _index: doc.index, _id: doc.id }
    }
  },
  transformDocument (doc) {
    return doc.body
  }
})
```
Another solution, but it could cause a breaking change in some cases (which is not a big deal given that the helpers are still experimental):
```js
const { createReadStream } = require('fs')
const split = require('split2')
const { Client } = require('@elastic/elasticsearch')

// assuming the document has the following shape:
// { index: string, id: string, body: object }

const client = new Client({ node: 'http://localhost:9200' })
const result = await client.helpers.bulk({
  datasource: createReadStream('./dataset.ndjson').pipe(split(JSON.parse)),
  onDocument (doc) {
    return {
      index: { _index: doc.index, _id: doc.id },
      document: doc.body
    }
  }
})
```

Pull request:
Allow document to be overwritten in `onDocument` iteratee of bulk helper
Allows the document to be submitted to be altered in the `onDocument` iteratee of the `bulk` method of the Elasticsearch Client Helpers.

Fixes #1209

Patch:
diff --git a/src/helpers.ts b/src/helpers.ts
index 9c5c822b..768ad1dc 100644
--- a/src/helpers.ts
+++ b/src/helpers.ts
@@ -74,11 +74,11 @@ export interface BulkStats {
   aborted: boolean
 }
 
-interface IndexAction {
+interface IndexActionOperation {
   index: T.BulkIndexOperation
 }
 
-interface CreateAction {
+interface CreateActionOperation {
   create: T.BulkCreateOperation
 }
 
@@ -90,7 +90,9 @@ interface DeleteAction {
   delete: T.BulkDeleteOperation
 }
 
-type UpdateAction = [UpdateActionOperation, Record<string, any>]
+type CreateAction = CreateActionOperation | [CreateActionOperation, unknown]
+type IndexAction = IndexActionOperation | [IndexActionOperation, unknown]
+type UpdateAction = [UpdateActionOperation, T.BulkUpdateAction]
 type Action = IndexAction | CreateAction | UpdateAction | DeleteAction
 
 export interface OnDropDocument<TDocument = unknown> {
@@ -618,22 +620,21 @@ export default class Helpers {
       for await (const chunk of datasource) {
         if (shouldAbort) break
         timeoutRef.refresh()
-        const action = onDocument(chunk)
-        const operation = Array.isArray(action)
-          ? Object.keys(action[0])[0]
-          : Object.keys(action)[0]
+        const result = onDocument(chunk)
+        const [action, payload] = Array.isArray(result) ? result : [result, chunk]
+        const operation = Object.keys(action)[0]
         if (operation === 'index' || operation === 'create') {
           actionBody = serializer.serialize(action)
-          payloadBody = typeof chunk === 'string' ? chunk : serializer.serialize(chunk)
+          payloadBody = typeof payload === 'string'
+            ? payload
+            : serializer.serialize(payload)
           chunkBytes += Buffer.byteLength(actionBody) + Buffer.byteLength(payloadBody)
           bulkBody.push(actionBody, payloadBody)
         } else if (operation === 'update') {
-          // @ts-expect-error in case of update action is an array
-          actionBody = serializer.serialize(action[0])
+          actionBody = serializer.serialize(action)
           payloadBody = typeof chunk === 'string'
             ? `{"doc":${chunk}}`
-            // @ts-expect-error in case of update action is an array
-            : serializer.serialize({ doc: chunk, ...action[1] })
+            : serializer.serialize({ doc: chunk, ...payload })
           chunkBytes += Buffer.byteLength(actionBody) + Buffer.byteLength(payloadBody)
           bulkBody.push(actionBody, payloadBody)
         } else if (operation === 'delete') {


Test patch:
diff --git a/test/unit/helpers/bulk.test.ts b/test/unit/helpers/bulk.test.ts
index dbabef07..732d696c 100644
--- a/test/unit/helpers/bulk.test.ts
+++ b/test/unit/helpers/bulk.test.ts
@@ -17,11 +17,11 @@
  * under the License.
  */
 
-import * as http from 'http'
+import FakeTimers from '@sinonjs/fake-timers'
 import { createReadStream } from 'fs'
+import * as http from 'http'
 import { join } from 'path'
 import split from 'split2'
-import FakeTimers from '@sinonjs/fake-timers'
 import { test } from 'tap'
 import { Client, errors } from '../../../'
 import { buildServer, connection } from '../../utils'
@@ -785,6 +785,59 @@ test('bulk index', t => {
     t.end()
   })
 
+  t.test('Should use payload returned by `onDocument`', async t => {
+    let count = 0
+    const updatedAt = '1970-01-01T12:00:00.000Z'
+    const MockConnection = connection.buildMockConnection({
+      onRequest (params) {
+        t.equal(params.path, '/_bulk')
+        t.match(params.headers, {
+          'content-type': 'application/vnd.elasticsearch+x-ndjson; compatible-with=8',
+          'x-elastic-client-meta': `es=${clientVersion},js=${nodeVersion},t=${transportVersion},hc=${nodeVersion},h=bp`
+        })
+        // @ts-expect-error
+        const [action, payload] = params.body.split('\n')
+        t.same(JSON.parse(action), { index: { _index: 'test' } })
+        t.same(JSON.parse(payload), { ...dataset[count++], updatedAt })
+        return { body: { errors: false, items: [{}] } }
+      }
+    })
+
+    const client = new Client({
+      node: 'http://localhost:9200',
+      Connection: MockConnection
+    })
+    const result = await client.helpers.bulk<Document>({
+      datasource: dataset.slice(),
+      flushBytes: 1,
+      concurrency: 1,
+      onDocument (doc) {
+        t.type(doc.user, 'string') // testing that doc is type of Document
+        return [
+          { 
+            index: { 
+              _index: 'test' 
+            } 
+          }, 
+          { ...doc, updatedAt }
+        ]
+      },
+      onDrop (doc) {
+        t.fail('This should never be called')
+      }
+    })
+
+    t.type(result.time, 'number')
+    t.type(result.bytes, 'number')
+    t.match(result, {
+      total: 3,
+      successful: 3,
+      retry: 0,
+      failed: 0,
+      aborted: false
+    })
+  })
+
   t.end()
 })
 
@@ -835,6 +888,58 @@ test('bulk create', t => {
       aborted: false
     })
   })
+
+  t.test('Should use payload returned by `onDocument`', async t => {
+    let count = 0
+    const updatedAt = '1970-01-01T12:00:00.000Z'
+    const MockConnection = connection.buildMockConnection({
+      onRequest (params) {
+        t.equal(params.path, '/_bulk')
+        t.match(params.headers, { 'content-type': 'application/vnd.elasticsearch+x-ndjson; compatible-with=8' })
+        // @ts-expect-error
+        const [action, payload] = params.body.split('\n')
+        t.same(JSON.parse(action), { create: { _index: 'test', _id: count } })
+        t.same(JSON.parse(payload), { ...dataset[count++], updatedAt })
+        return { body: { errors: false, items: [{}] } }
+      }
+    })
+
+    const client = new Client({
+      node: 'http://localhost:9200',
+      Connection: MockConnection
+    })
+    let id = 0
+    const result = await client.helpers.bulk({
+      datasource: dataset.slice(),
+      flushBytes: 1,
+      concurrency: 1,
+      onDocument (doc) {
+        return [
+          {
+            create: {
+              _index: 'test',
+              _id: String(id++)
+            }
+          }, 
+          { ...doc, updatedAt }
+        ]
+      },
+      onDrop (doc) {
+        t.fail('This should never be called')
+      }
+    })
+
+    t.type(result.time, 'number')
+    t.type(result.bytes, 'number')
+    t.match(result, {
+      total: 3,
+      successful: 3,
+      retry: 0,
+      failed: 0,
+      aborted: false
+    })
+  })
+
   t.end()
 })