Add Document Key Retriever to Buffered Sender (Azure#13478)

* Add Document Key Retriever to Buffered Sender

* Fix Lint Errors

* Update sdk/search/search-documents/src/searchIndexingBufferedSender.ts

Co-authored-by: Jeff Fisher <[email protected]>

* Minor changes

* Remove ts-ignore comment

Co-authored-by: Jeff Fisher <[email protected]>

Author: sarangan12

sdk/search/search-documents/review/search-documents.api.md

@@ -1653,7 +1653,7 @@ export interface SearchIndexerWarning {
 
 // @public
 export class SearchIndexingBufferedSender<T> {
-    constructor(client: IndexDocumentsClient<T>, options?: SearchIndexingBufferedSenderOptions);
+    constructor(client: IndexDocumentsClient<T>, documentKeyRetriever: (document: T) => string, options?: SearchIndexingBufferedSenderOptions);
     deleteDocuments(documents: T[], options?: SearchIndexingBufferedSenderDeleteDocumentsOptions): Promise<void>;
     dispose(): Promise<void>;
     flush(options?: SearchIndexingBufferedSenderFlushDocumentsOptions): Promise<void>;

sdk/search/search-documents/samples/typescript/src/bufferedSender/uploadDocuments/autoFlushSizeBased.ts

@@ -5,7 +5,7 @@ import {
   GeographyPoint,
   SearchIndexClient
 } from "@azure/search-documents";
-import { createIndex, WAIT_TIME } from "../../utils/setup";
+import { createIndex, documentKeyRetriever, WAIT_TIME } from "../../utils/setup";
 import { Hotel } from "../../utils/interfaces";
 import { delay } from "@azure/core-http";
 import * as dotenv from "dotenv";
@@ -65,9 +65,13 @@ export async function main() {
   await createIndex(indexClient, TEST_INDEX_NAME);
   await delay(WAIT_TIME);
 
-  const bufferedClient = new SearchIndexingBufferedSender<Hotel>(searchClient, {
-    autoFlush: true
-  });
+  const bufferedClient = new SearchIndexingBufferedSender<Hotel>(
+    searchClient,
+    documentKeyRetriever,
+    {
+      autoFlush: true
+    }
+  );
 
   bufferedClient.on("batchAdded", (response: any) => {
     console.log(`Batch Added Event has been receieved: ${response}`);

sdk/search/search-documents/samples/typescript/src/bufferedSender/uploadDocuments/autoFlushTimerBased.ts

@@ -6,7 +6,7 @@ import {
   SearchIndexClient,
   DEFAULT_FLUSH_WINDOW
 } from "@azure/search-documents";
-import { createIndex, WAIT_TIME } from "../../utils/setup";
+import { createIndex, documentKeyRetriever, WAIT_TIME } from "../../utils/setup";
 import { Hotel } from "../../utils/interfaces";
 import { delay } from "@azure/core-http";
 import * as dotenv from "dotenv";
@@ -39,6 +39,7 @@ export async function main() {
 
   const bufferedClient: SearchIndexingBufferedSender<Hotel> = new SearchIndexingBufferedSender(
     searchClient,
+    documentKeyRetriever,
     {
       autoFlush: true
     }

sdk/search/search-documents/samples/typescript/src/bufferedSender/uploadDocuments/manualFlush.ts

@@ -5,7 +5,7 @@ import {
   GeographyPoint,
   SearchIndexClient
 } from "@azure/search-documents";
-import { createIndex, WAIT_TIME } from "../../utils/setup";
+import { createIndex, documentKeyRetriever, WAIT_TIME } from "../../utils/setup";
 import { Hotel } from "../../utils/interfaces";
 import { delay } from "@azure/core-http";
 import * as dotenv from "dotenv";
@@ -36,6 +36,7 @@ export async function main() {
 
   const bufferedClient: SearchIndexingBufferedSender<Hotel> = new SearchIndexingBufferedSender(
     searchClient,
+    documentKeyRetriever,
     {
       autoFlush: false
     }

sdk/search/search-documents/samples/typescript/src/utils/setup.ts

@@ -2,9 +2,14 @@
 // Licensed under the MIT license.
 
 import { SearchIndexClient, SearchIndex, KnownAnalyzerNames } from "@azure/search-documents";
+import { Hotel } from "./interfaces";
 
 export const WAIT_TIME = 4000;
 
+export const documentKeyRetriever: (document: Hotel) => string = (document: Hotel): string => {
+  return document.hotelId;
+};
+
 // eslint-disable-next-line @azure/azure-sdk/ts-use-interface-parameters
 export async function createIndex(client: SearchIndexClient, name: string): Promise<void> {
   const hotelIndex: SearchIndex = {

sdk/search/search-documents/src/generated/data/searchClient.ts

@@ -10,7 +10,7 @@ import { Documents } from "./operations";
 import { SearchClientContext } from "./searchClientContext";
 import { SearchClientOptionalParams, ApiVersion20200630 } from "./models";
 
-/** @hidden */
+/** @internal */
 export class SearchClient extends SearchClientContext {
   /**
    * Initializes a new instance of the SearchClient class.

sdk/search/search-documents/src/generated/data/searchClientContext.ts

@@ -12,7 +12,7 @@ import { ApiVersion20200630, SearchClientOptionalParams } from "./models";
 const packageName = "@azure/search-documents";
 const packageVersion = "11.1.0-beta.2";
 
-/** @hidden */
+/** @internal */
 export class SearchClientContext extends coreHttp.ServiceClient {
   endpoint: string;
   indexName: string;

sdk/search/search-documents/src/generated/service/searchServiceClient.ts

@@ -24,7 +24,7 @@ import {
   SearchServiceClientGetServiceStatisticsResponse
 } from "./models";
 
-/** @hidden */
+/** @internal */
 export class SearchServiceClient extends SearchServiceClientContext {
   /**
    * Initializes a new instance of the SearchServiceClient class.

sdk/search/search-documents/src/generated/service/searchServiceClientContext.ts

@@ -15,7 +15,7 @@ import {
 const packageName = "@azure/search-documents";
 const packageVersion = "11.1.0-beta.2";
 
-/** @hidden */
+/** @internal */
 export class SearchServiceClientContext extends coreHttp.ServiceClient {
   endpoint: string;
   apiVersion: ApiVersion20200630;

sdk/search/search-documents/src/searchIndexingBufferedSender.ts

@@ -103,6 +103,10 @@ export class SearchIndexingBufferedSender<T> {
    * Event emitter/publisher used in the Buffered Sender
    */
   private readonly emitter = new EventEmitter();
+  /**
+   * Method to retrieve the document key
+   */
+  private documentKeyRetriever: (document: T) => string;
 
   /**
    * Creates a new instance of SearchIndexingBufferedSender.
@@ -111,10 +115,15 @@ export class SearchIndexingBufferedSender<T> {
    * @param options - Options to modify auto flush.
    *
    */
-  constructor(client: IndexDocumentsClient<T>, options: SearchIndexingBufferedSenderOptions = {}) {
+  constructor(
+    client: IndexDocumentsClient<T>,
+    documentKeyRetriever: (document: T) => string,
+    options: SearchIndexingBufferedSenderOptions = {}
+  ) {
     this.client = client;
+    this.documentKeyRetriever = documentKeyRetriever;
     // General Configuration properties
-    this.autoFlush = options.autoFlush ?? false;
+    this.autoFlush = options.autoFlush ?? true;
     this.initialBatchActionCount = options.initialBatchActionCount ?? DEFAULT_BATCH_SIZE;
     this.flushWindowInMs = options.flushWindowInMs ?? DEFAULT_FLUSH_WINDOW;
     // Retry specific configuration properties
@@ -377,9 +386,30 @@ export class SearchIndexingBufferedSender<T> {
       this.batchObject = new IndexDocumentsBatch<T>();
       while (actions.length > 0) {
         const actionsToSend = actions.splice(0, this.initialBatchActionCount);
-        await this.submitDocuments(actionsToSend, options);
+        const { batchToSubmit, submitLater } = this.pruneActions(actionsToSend);
+        actions.unshift(...submitLater);
+        await this.submitDocuments(batchToSubmit, options);
+      }
+    }
+  }
+
+  private pruneActions(
+    batch: IndexDocumentsAction<T>[]
+  ): { batchToSubmit: IndexDocumentsAction<T>[]; submitLater: IndexDocumentsAction<T>[] } {
+    const hashSet: Set<string> = new Set<string>();
+    const resultBatch: IndexDocumentsAction<T>[] = [];
+    const pruned: IndexDocumentsAction<T>[] = [];
+
+    for (const document of batch) {
+      const key = this.documentKeyRetriever((document as unknown) as T);
+      if (hashSet.has(key)) {
+        pruned.push(document);
+      } else {
+        hashSet.add(key);
+        resultBatch.push(document);
       }
     }
+    return { batchToSubmit: resultBatch, submitLater: pruned };
   }
 
   private async submitDocuments(
@@ -398,7 +428,7 @@ export class SearchIndexingBufferedSender<T> {
       // raise success event
       this.emitter.emit("batchSucceeded", result);
     } catch (e) {
-      if (e.code && e.code === "413" && actionsToSend.length > 1) {
+      if (e.statusCode && e.statusCode === 413 && actionsToSend.length > 1) {
         // Cut the payload size to half
         const splitActionsArray = [
           actionsToSend.slice(0, actionsToSend.length / 2),
@@ -427,6 +457,6 @@ export class SearchIndexingBufferedSender<T> {
   }
 
   private isRetryAbleError(e: any): boolean {
-    return e.code && (e.code === "422" || e.code === "409" || e.code === "503");
+    return e.statusCode && (e.statusCode === 422 || e.statusCode === 409 || e.statusCode === 503);
   }
 }