Import the latest main branch.

Signed-off-by: huuyafwww <huuya1234fwww@gmail.com>
opensearch-project · Apr 15, 2023 · 2e47c79 · 2e47c79
1 parent 1951913
commit 2e47c79
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 9 deletions.
diff --git a/lib/Helpers.d.ts b/lib/Helpers.d.ts
@@ -102,6 +102,18 @@ interface UpdateActionOperation {
   };
 }
 
+interface UpdateActionDocOperation {
+  doc: {
+    [key: string]: any;
+  };
+}
+
+interface UpdateActionScriptOperation {
+  script: {
+    [key: string]: any;
+  };
+}
+
 interface DeleteAction {
   delete: {
     _index: string;
@@ -111,7 +123,10 @@ interface DeleteAction {
 
 type CreateAction = CreateActionOperation | [CreateActionOperation, unknown];
 type IndexAction = IndexActionOperation | [IndexActionOperation, unknown];
-type UpdateAction = [UpdateActionOperation, Record<string, any>];
+type UpdateAction = [
+  UpdateActionOperation,
+  UpdateActionDocOperation | UpdateActionScriptOperation
+];
 type Action = IndexAction | CreateAction | UpdateAction | DeleteAction;
 type Omit<T, K extends keyof T> = Pick<T, Exclude<keyof T, K>>;
 

diff --git a/lib/Helpers.js b/lib/Helpers.js
@@ -539,10 +539,7 @@ class Helpers {
           bulkBody.push(actionBody, payloadBody);
         } else if (operation === 'update') {
           actionBody = serializer.serialize(action);
-          payloadBody =
-            typeof chunk === 'string'
-              ? `{"doc":${chunk}}`
-              : serializer.serialize({ doc: chunk, ...payload });
+          payloadBody = serializer.serialize(payload || chunk);
           chunkBytes += Buffer.byteLength(actionBody) + Buffer.byteLength(payloadBody);
           bulkBody.push(actionBody, payloadBody);
         } else if (operation === 'delete') {

diff --git a/test/types/helpers.test-d.ts b/test/types/helpers.test-d.ts
@@ -36,6 +36,7 @@ import {
   ScrollSearchResponse,
   OnDropDocument,
   MsearchHelper,
+  UpdateActionDocOperation,
 } from '../../lib/Helpers';
 import { ApiResponse, ApiError, Context } from '../../lib/Transport';
 
@@ -107,7 +108,7 @@ expectError(
   const options: BulkHelperOptions<Record<string, any>> = {
     datasource: [],
     onDocument(doc: Record<string, any>) {
-      return [{ update: { _index: 'test' } }, doc];
+      return [{ update: { _index: 'test' } }, doc as UpdateActionDocOperation];
     },
   };
   expectAssignable<BulkHelperOptions<Record<string, any>>>(options);

diff --git a/test/unit/helpers/bulk.test.js b/test/unit/helpers/bulk.test.js
@@ -971,14 +971,16 @@ test('bulk update', (t) => {
       flushBytes: 1,
       concurrency: 1,
       onDocument() {
+        const currentId = id++;
         return [
           {
             update: {
               _index: 'test',
-              _id: id++,
+              _id: currentId,
             },
           },
           {
+            doc: dataset[currentId],
             doc_as_upsert: true,
           },
         ];
@@ -1022,13 +1024,15 @@ test('bulk update', (t) => {
       flushBytes: 1,
       concurrency: 1,
       onDocument() {
+        const currentId = id++;
         return [
           {
             update: {
               _index: 'test',
-              _id: id++,
+              _id: currentId,
             },
           },
+          { doc: dataset[currentId] },
         ];
       },
       onDrop() {
@@ -1070,14 +1074,16 @@ test('bulk update', (t) => {
       flushBytes: 1,
       concurrency: 1,
       onDocument() {
+        const currentId = id++;
         return [
           {
             update: {
               _index: 'test',
-              _id: id++,
+              _id: currentId,
             },
           },
           {
+            doc: dataset[currentId],
             doc_as_upsert: true,
           },
         ];