Fix rejected patches from elastic/elasticsearch-js#2199

Author: JoshMock

src/helpers.ts

@@ -128,6 +128,7 @@ export interface BulkHelperOptions<TDocument = unknown> extends T.BulkRequest {
   retries?: number
   wait?: number
   onDrop?: (doc: OnDropDocument<TDocument>) => void
+  onSuccess?: (doc: OnSuccessDocument) => void
 }
 
 export interface BulkHelper<T> extends Promise<BulkStats> {
@@ -397,7 +398,7 @@ export default class Helpers {
         clearTimeout(timeoutRef)
       }
 
-      // In some cases the previos http call does not have finished,
+      // In some cases the previous http call does not have finished,
       // or we didn't reach the flush bytes threshold, so we force one last operation.
       if (loadedOperations > 0) {
         const send = await semaphore()
@@ -433,8 +434,8 @@ export default class Helpers {
     // to guarantee that no more than the number of operations
     // allowed to run at the same time are executed.
     // It returns a semaphore function which resolves in the next tick
-    // if we didn't reach the maximim concurrency yet, otherwise it returns
-    // a promise that resolves as soon as one of the running request has finshed.
+    // if we didn't reach the maximum concurrency yet, otherwise it returns
+    // a promise that resolves as soon as one of the running requests has finished.
     // The semaphore function resolves a send function, which will be used
     // to send the actual msearch request.
     // It also returns a finish function, which returns a promise that is resolved
@@ -566,6 +567,9 @@ export default class Helpers {
       retries = this[kMaxRetries],
       wait = 5000,
       onDrop = noop,
+      // onSuccess does not default to noop, to avoid the performance hit
+      // of deserializing every document in the bulk request
+      onSuccess,
       ...bulkOptions
     } = options
 
@@ -638,7 +642,7 @@ export default class Helpers {
       let chunkBytes = 0
       timeoutRef = setTimeout(onFlushTimeout, flushInterval) // eslint-disable-line
 
-      // @ts-expect-error datasoruce is an iterable
+      // @ts-expect-error datasource is an iterable
       for await (const chunk of datasource) {
         if (shouldAbort) break
         timeoutRef.refresh()
@@ -679,7 +683,7 @@ export default class Helpers {
       }
 
       clearTimeout(timeoutRef)
-      // In some cases the previos http call does not have finished,
+      // In some cases the previous http call has not finished,
       // or we didn't reach the flush bytes threshold, so we force one last operation.
       if (!shouldAbort && chunkBytes > 0) {
         const send = await semaphore()
@@ -715,8 +719,8 @@ export default class Helpers {
     // to guarantee that no more than the number of operations
     // allowed to run at the same time are executed.
     // It returns a semaphore function which resolves in the next tick
-    // if we didn't reach the maximim concurrency yet, otherwise it returns
-    // a promise that resolves as soon as one of the running request has finshed.
+    // if we didn't reach the maximum concurrency yet, otherwise it returns
+    // a promise that resolves as soon as one of the running requests has finished.
     // The semaphore function resolves a send function, which will be used
     // to send the actual bulk request.
     // It also returns a finish function, which returns a promise that is resolved
@@ -823,57 +827,93 @@ export default class Helpers {
         callback()
       }
 
+      /**
+       * Zips bulk response items (the action's result) with the original document body.
+       * The raw string version of action and document lines are also included.
+       */
+      function zipBulkResults (responseItems: BulkResponseItem[], bulkBody: string[]): ZippedResult[] {
+        const zipped = []
+        let indexSlice = 0
+        for (let i = 0, len = responseItems.length; i < len; i++) {
+          const result = responseItems[i]
+          const operation = Object.keys(result)[0]
+          let zipResult
+
+          if (operation === 'delete') {
+            zipResult = {
+              result,
+              raw: { action: bulkBody[indexSlice] }
+            }
+            indexSlice += 1
+          } else {
+            const document = bulkBody[indexSlice + 1]
+            zipResult = {
+              result,
+              raw: { action: bulkBody[indexSlice], document },
+              // this is a function so that deserialization is only done when needed
+              // to avoid a performance hit
+              document: () => serializer.deserialize(document)
+            }
+            indexSlice += 2
+          }
+
+          zipped.push(zipResult as ZippedResult)
+        }
+
+        return zipped
+      }
+
       function tryBulk (bulkBody: string[], callback: (err: Error | null, bulkBody: string[]) => void): void {
         if (shouldAbort) return callback(null, [])
         client.bulk(Object.assign({}, bulkOptions, { body: bulkBody }), reqOptions as TransportRequestOptionsWithMeta)
           .then(response => {
             const result = response.body
+            const results = zipBulkResults(result.items, bulkBody)
+
             if (!result.errors) {
               stats.successful += result.items.length
-              for (const item of result.items) {
-                if (item.update?.result === 'noop') {
+              for (const item of results) {
+                const { result, document = noop } = item
+                if (result.update?.result === 'noop') {
                   stats.noop++
                 }
+                if (onSuccess != null) onSuccess({ result, document: document() })
               }
               return callback(null, [])
             }
             const retry = []
-            const { items } = result
-            let indexSlice = 0
-            for (let i = 0, len = items.length; i < len; i++) {
-              const action = items[i]
-              const operation = Object.keys(action)[0]
+            for (const item of results) {
+              const { result, raw, document = noop } = item
+              const operation = Object.keys(result)[0]
               // @ts-expect-error
-              const responseItem = action[operation as keyof T.BulkResponseItemContainer]
+              const responseItem = result[operation as keyof T.BulkResponseItemContainer]
               assert(responseItem !== undefined, 'The responseItem is undefined, please file a bug report')
 
               if (responseItem.status >= 400) {
                 // 429 is the only status code where we might want to retry
                 // a document, because it was not an error in the document itself,
-                // but the ES node were handling too many operations.
+                // but the ES node was handling too many operations.
                 if (responseItem.status === 429) {
-                  retry.push(bulkBody[indexSlice])
+                  retry.push(raw.action)
                   /* istanbul ignore next */
                   if (operation !== 'delete') {
-                    retry.push(bulkBody[indexSlice + 1])
+                    retry.push(raw.document ?? '')
                   }
                 } else {
                   onDrop({
                     status: responseItem.status,
                     error: responseItem.error ?? null,
-                    operation: serializer.deserialize(bulkBody[indexSlice]),
+                    operation: serializer.deserialize(raw.action),
                     // @ts-expect-error
-                    document: operation !== 'delete'
-                      ? serializer.deserialize(bulkBody[indexSlice + 1])
-                      : null,
+                    document: document(),
                     retried: isRetrying
                   })
                   stats.failed += 1
                 }
               } else {
                 stats.successful += 1
+                if (onSuccess != null) onSuccess({ result, document: document() })
               }
-              operation === 'delete' ? indexSlice += 1 : indexSlice += 2
             }
             callback(null, retry)
           })

test/unit/helpers/bulk.test.ts

@@ -17,12 +17,12 @@
  * under the License.
  */
 
+import FakeTimers from '@sinonjs/fake-timers'
 import { AssertionError } from 'assert'
-import * as http from 'http'
 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'
@@ -610,7 +610,8 @@ test('bulk index', t => {
 
       const client = new Client({
         node: 'http://localhost:9200',
-        Connection: MockConnection
+        Connection: MockConnection,
+        compression: false
       })
 
       let count = 0
@@ -701,7 +702,8 @@ test('bulk index', t => {
 
       const client = new Client({
         node: 'http://localhost:9200',
-        Connection: MockConnection
+        Connection: MockConnection,
+        compression: false
       })
       const stream = createReadStream(join(__dirname, '..', '..', 'fixtures', 'small-dataset.ndjson'), 'utf8')
 
@@ -795,7 +797,8 @@ test('bulk index', t => {
 
       const client = new Client({
         node: 'http://localhost:9200',
-        Connection: MockConnection
+        Connection: MockConnection,
+        compression: false
       })
 
       async function * generator () {
@@ -1033,8 +1036,6 @@ test('bulk update', t => {
     })
   })
 
-
-
   t.end()
 })
 
@@ -1164,7 +1165,7 @@ test('bulk delete', t => {
   t.test('Should call onDrop on the correct document when doing a mix of operations that includes deletes', async t => {
     // checks to ensure onDrop doesn't provide the wrong document when some operations are deletes
     // see https://github.com/elastic/elasticsearch-js/issues/1751
-    async function handler (req: http.IncomingMessage, res: http.ServerResponse) {
+    async function handler (_req: http.IncomingMessage, res: http.ServerResponse) {
       res.setHeader('content-type', 'application/json')
       res.end(JSON.stringify({
         took: 0,
@@ -1178,14 +1179,17 @@ test('bulk delete', t => {
     }
 
     const [{ port }, server] = await buildServer(handler)
-    const client = new Client({ node: `http://localhost:${port}` })
+    const client = new Client({
+      node: `http://localhost:${port}`,
+      compression: false
+    })
     let counter = 0
     const result = await client.helpers.bulk({
       datasource: dataset.slice(),
       concurrency: 1,
       wait: 10,
       retries: 0,
-      onDocument (doc) {
+      onDocument (_doc) {
         counter++
         if (counter === 1) {
           return {
@@ -1237,7 +1241,8 @@ test('bulk delete', t => {
 
     const client = new Client({
       node: 'http://localhost:9200',
-      Connection: MockConnection
+      Connection: MockConnection,
+      compression: false
     })
 
     let docCount = 0
@@ -1464,12 +1469,12 @@ test('Flush interval', t => {
       })(),
       flushBytes: 5000000,
       concurrency: 1,
-      onDocument (doc) {
+      onDocument (_doc) {
         return {
           index: { _index: 'test' }
         }
       },
-      onDrop (doc) {
+      onDrop (_doc) {
         t.fail('This should never be called')
       }
     })
@@ -1525,12 +1530,12 @@ test('Flush interval', t => {
       })(),
       flushBytes: 5000000,
       concurrency: 1,
-      onDocument (doc) {
+      onDocument (_doc) {
         return {
           index: { _index: 'test' }
         }
       },
-      onDrop (doc) {
+      onDrop (_doc) {
         t.fail('This should never be called')
       }
     })