returnDocumentResponses implementation (#67)

toptobes · web-flow · commit 10362d881f1c · 2024-08-15T15:21:51.000-05:00
* Ok

* returnDocumentResponses implementation + partial docs
diff --git a/src/data-api/collection.ts b/src/data-api/collection.ts
@@ -14,13 +14,12 @@
 
 import { normalizeSort } from './utils';
 import { FindCursor } from '@/src/data-api/cursor';
-import { Db, SomeDoc, SomeId } from '@/src/data-api';
+import { DataAPIDetailedErrorDescriptor, Db, SomeDoc, SomeId } from '@/src/data-api';
 import {
   BulkWriteError,
   CollectionNotFoundError,
   DataAPIResponseError,
-  DeleteManyError,
-  InsertManyError,
+  DeleteManyError, InsertManyError,
   mkRespErrorFromResponse,
   mkRespErrorFromResponses,
   TooManyDocumentsToCountError,
@@ -73,7 +72,7 @@ import { FindOneAndReplaceCommand } from '@/src/data-api/types/find/find-one-rep
 import { DeleteOneCommand } from '@/src/data-api/types/delete/delete-one';
 import { FindOneAndDeleteCommand } from '@/src/data-api/types/find/find-one-delete';
 import { FindOneAndUpdateCommand } from '@/src/data-api/types/find/find-one-update';
-import { InsertManyCommand } from '@/src/data-api/types/insert/insert-many';
+import { InsertManyCommand, InsertManyDocumentResponse } from '@/src/data-api/types/insert/insert-many';
 import { Mutable } from '@/src/data-api/types/utils';
 import { CollectionSpawnOptions } from '@/src/data-api/types/collections/spawn-collection';
 
@@ -474,7 +473,7 @@ export class Collection<Schema extends SomeDoc = SomeDoc> {
       commonResult.matchedCount += desc.rawResponse.status?.matchedCount ?? 0;
       commonResult.upsertedCount = desc.rawResponse.status?.upsertedCount ?? 0;
 
-      throw mkRespErrorFromResponse(UpdateManyError, command, desc.rawResponse, commonResult);
+      throw mkRespErrorFromResponse(UpdateManyError, command, desc.rawResponse, { partialResult: commonResult });
     }
 
     return (resp.status?.upsertedId)
@@ -677,8 +676,14 @@ export class Collection<Schema extends SomeDoc = SomeDoc> {
       if (!(e instanceof DataAPIResponseError)) {
         throw e;
       }
+
       const desc = e.detailedErrorDescriptors[0];
-      throw mkRespErrorFromResponse(DeleteManyError, command, desc.rawResponse, { deletedCount: numDeleted + (desc.rawResponse.status?.deletedCount ?? 0) })
+
+      throw mkRespErrorFromResponse(DeleteManyError, command, desc.rawResponse, {
+        partialResult: {
+          deletedCount: numDeleted + (desc.rawResponse.status?.deletedCount ?? 0),
+        },
+      });
     }
 
     return {
@@ -1516,35 +1521,37 @@ const coalesceVectorSpecialsIntoSort = <T extends OptionsWithSort>(options: T |
 
 // -- Insert Many ------------------------------------------------------------------------------------------
 
-const insertManyOrdered = async <Schema>(httpClient: DataAPIHttpClient, documents: unknown[], chunkSize: number, timeoutManager: TimeoutManager): Promise<IdOf<Schema>[]> => {
+const insertManyOrdered = async <Schema extends SomeDoc>(httpClient: DataAPIHttpClient, documents: unknown[], chunkSize: number, timeoutManager: TimeoutManager): Promise<IdOf<Schema>[]> => {
   const insertedIds: IdOf<Schema>[] = [];
 
   for (let i = 0, n = documents.length; i < n; i += chunkSize) {
     const slice = documents.slice(i, i + chunkSize);
 
-    try {
-      const inserted = await insertMany<Schema>(httpClient, slice, true, timeoutManager);
-      insertedIds.push(...inserted);
-    } catch (e) {
-      if (!(e instanceof DataAPIResponseError)) {
-        throw e;
-      }
-      const desc = e.detailedErrorDescriptors[0];
+    const [docResp, inserted, errDesc] = await insertMany<Schema>(httpClient, slice, true, timeoutManager);
+    insertedIds.push(...inserted);
 
-      insertedIds.push(...desc.rawResponse.status?.insertedIds ?? []);
-      throw mkRespErrorFromResponse(InsertManyError, desc.command, desc.rawResponse, { insertedIds: insertedIds as SomeId[], insertedCount: insertedIds.length })
+    if (errDesc) {
+      throw mkRespErrorFromResponse(InsertManyError, errDesc.command, errDesc.rawResponse, {
+        partialResult: {
+          insertedIds: insertedIds as SomeId[],
+          insertedCount: insertedIds.length,
+        },
+        documentResponses: docResp,
+        failedCount: docResp.length - insertedIds.length,
+      });
     }
   }
 
   return insertedIds;
 }
 
-const insertManyUnordered = async <Schema>(httpClient: DataAPIHttpClient, documents: unknown[], concurrency: number, chunkSize: number, timeoutManager: TimeoutManager): Promise<IdOf<Schema>[]> => {
+const insertManyUnordered = async <Schema extends SomeDoc>(httpClient: DataAPIHttpClient, documents: unknown[], concurrency: number, chunkSize: number, timeoutManager: TimeoutManager): Promise<IdOf<Schema>[]> => {
   const insertedIds: IdOf<Schema>[] = [];
   let masterIndex = 0;
 
   const failCommands = [] as Record<string, any>[];
   const failRaw = [] as Record<string, any>[];
+  const docResps = [] as InsertManyDocumentResponse<SomeDoc>[];
 
   const workers = Array.from({ length: concurrency }, async () => {
     while (masterIndex < documents.length) {
@@ -1558,42 +1565,72 @@ const insertManyUnordered = async <Schema>(httpClient: DataAPIHttpClient, docume
 
       const slice = documents.slice(localI, endIdx);
 
-      try {
-        const inserted = await insertMany<Schema>(httpClient, slice, false, timeoutManager);
-        insertedIds.push(...inserted);
-      } catch (e) {
-        if (!(e instanceof DataAPIResponseError)) {
-          throw e;
-        }
-        const desc = e.detailedErrorDescriptors[0];
-
-        const justInserted = desc.rawResponse.status?.insertedIds ?? [];
-        insertedIds.push(...justInserted);
+      const [docResp, inserted, errDesc] = await insertMany<Schema>(httpClient, slice, false, timeoutManager);
+      insertedIds.push(...inserted);
+      docResps.push(...docResp);
 
-        failCommands.push(desc.command);
-        failRaw.push(desc.rawResponse);
+      if (errDesc) {
+        failCommands.push(errDesc.command);
+        failRaw.push(errDesc.rawResponse);
       }
     }
   });
   await Promise.all(workers);
 
   if (failCommands.length > 0) {
-    throw mkRespErrorFromResponses(InsertManyError, failCommands, failRaw, { insertedIds: insertedIds as SomeId[], insertedCount: insertedIds.length });
+    throw mkRespErrorFromResponses(InsertManyError, failCommands, failRaw, {
+      partialResult: {
+        insertedIds: insertedIds as SomeId[],
+        insertedCount: insertedIds.length,
+      },
+      documentResponses: docResps,
+      failedCount: docResps.length - insertedIds.length,
+    });
   }
 
   return insertedIds;
 }
 
-const insertMany = async <Schema>(httpClient: DataAPIHttpClient, documents: unknown[], ordered: boolean, timeoutManager: TimeoutManager): Promise<IdOf<Schema>[]> => {
+const insertMany = async <Schema extends SomeDoc>(httpClient: DataAPIHttpClient, documents: unknown[], ordered: boolean, timeoutManager: TimeoutManager): Promise<[InsertManyDocumentResponse<Schema>[], IdOf<Schema>[], DataAPIDetailedErrorDescriptor | undefined]> => {
   const command: InsertManyCommand = {
     insertMany: {
       documents,
-      options: { ordered },
+      options: {
+        returnDocumentResponses: true,
+        ordered,
+      },
     }
   }
 
-  const resp = await httpClient.executeCommand(command, { timeoutManager });
-  return resp.status?.insertedIds ?? [];
+  let resp, err: DataAPIResponseError | undefined;
+
+  try {
+    resp = await httpClient.executeCommand(command, { timeoutManager });
+  } catch (e) {
+    if (! (e instanceof DataAPIResponseError)) {
+      throw e;
+    }
+    resp = e.detailedErrorDescriptors[0].rawResponse;
+    err = e;
+  }
+
+  const documentResponses = resp.status!.documentResponses;
+  const errors = resp.errors!;
+
+  const insertedIds = [];
+
+  for (let i = 0, n = documentResponses.length; i < n; i++) {
+    const resp = documentResponses[i];
+
+    if (resp.status === "OK") {
+      insertedIds.push(resp._id);
+    } else if (resp.errorIdx) {
+      resp.error = errors[resp.errorIdx];
+      delete resp.errorIdx;
+    }
+  }
+
+  return [documentResponses, insertedIds, err?.detailedErrorDescriptors[0]];
 }
 
 // -- Bulk Write ------------------------------------------------------------------------------------------
@@ -1616,7 +1653,7 @@ const bulkWriteOrdered = async <Schema extends SomeDoc>(httpClient: DataAPIHttpC
       addToBulkWriteResult(results, desc.rawResponse.status, i)
     }
 
-    throw mkRespErrorFromResponse(BulkWriteError, desc.command, desc.rawResponse, results);
+    throw mkRespErrorFromResponse(BulkWriteError, desc.command, desc.rawResponse, { partialResult: results });
   }
 
   return results;
@@ -1654,7 +1691,7 @@ const bulkWriteUnordered = async <Schema extends SomeDoc>(httpClient: DataAPIHtt
   await Promise.all(workers);
 
   if (failCommands.length > 0) {
-    throw mkRespErrorFromResponses(BulkWriteError, failCommands, failRaw, results);
+    throw mkRespErrorFromResponses(BulkWriteError, failCommands, failRaw, { partialResult: results });
   }
 
   return results;
diff --git a/src/data-api/errors.ts b/src/data-api/errors.ts
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import type { InsertManyResult } from '@/src/data-api/types/insert/insert-many';
+import type { InsertManyDocumentResponse, InsertManyResult } from '@/src/data-api/types/insert/insert-many';
 import type { DeleteManyResult } from '@/src/data-api/types/delete/delete-many';
 import type { UpdateManyResult } from '@/src/data-api/types/update/update-many';
 import type { BulkWriteResult } from '@/src/data-api/types/misc/bulk-write';
@@ -449,6 +449,19 @@ export class InsertManyError extends CumulativeDataAPIError {
    * of all successful insertions.
    */
   declare public readonly partialResult: InsertManyResult<SomeDoc>;
+
+  /**
+   * The specific statuses and ids for each document present in the `insertMany` command
+   *
+   * The position of each document response is the same as its corresponding document in the input `documents` array
+   */
+  declare public readonly documentResponses: InsertManyDocumentResponse<SomeDoc>[];
+
+  /**
+   * The number of documents which failed insertion (i.e. their status in {@link InsertManyError.documentResponses} was
+   * `'ERROR'` or `'SKIPPED'`)
+   */
+  declare public readonly failedCount: number;
 }
 
 /**
@@ -532,19 +545,17 @@ export class BulkWriteError extends CumulativeDataAPIError {
   declare public readonly partialResult: BulkWriteResult<SomeDoc>;
 }
 
-type InferPartialResult<T> = T extends { readonly partialResult: infer P } ? P : never;
-
 /**
  * @internal
  */
-export const mkRespErrorFromResponse = <E extends DataAPIResponseError>(err: new (message: string, errorDescriptors: DataAPIErrorDescriptor[], detailedErrorDescriptors: DataAPIDetailedErrorDescriptor[]) => E, command: Record<string, any>, raw: RawDataAPIResponse, partialResult?: InferPartialResult<E>) => {
-  return mkRespErrorFromResponses(err, [command], [raw], partialResult);
+export const mkRespErrorFromResponse = <E extends DataAPIResponseError>(err: new (message: string, errorDescriptors: DataAPIErrorDescriptor[], detailedErrorDescriptors: DataAPIDetailedErrorDescriptor[]) => E, command: Record<string, any>, raw: RawDataAPIResponse, attributes?: Omit<E, keyof DataAPIResponseError>) => {
+  return mkRespErrorFromResponses(err, [command], [raw], attributes);
 }
 
 /**
  * @internal
  */
-export const mkRespErrorFromResponses = <E extends DataAPIResponseError>(err: new (message: string, errorDescriptors: DataAPIErrorDescriptor[], detailedErrorDescriptors: DataAPIDetailedErrorDescriptor[]) => E, commands: Record<string, any>[], raw: RawDataAPIResponse[], partialResult?: InferPartialResult<E>) => {
+export const mkRespErrorFromResponses = <E extends DataAPIResponseError>(err: new (message: string, errorDescriptors: DataAPIErrorDescriptor[], detailedErrorDescriptors: DataAPIDetailedErrorDescriptor[]) => E, commands: Record<string, any>[], raw: RawDataAPIResponse[], attributes?: Omit<E, keyof DataAPIResponseError>) => {
   const detailedDescriptors = [] as DataAPIDetailedErrorDescriptor[];
 
   for (let i = 0, n = commands.length; i < n; i++) {
@@ -564,14 +575,9 @@ export const mkRespErrorFromResponses = <E extends DataAPIResponseError>(err: ne
   }
 
   const errorDescriptors = detailedDescriptors.flatMap(d => d.errorDescriptors);
-
   const message = errorDescriptors[0]?.message || 'Something unexpected occurred';
 
   const instance = new err(message, errorDescriptors, detailedDescriptors) ;
-
-  if (partialResult) {
-    /* @ts-expect-error - If the lord wants a partialResult, the lord will get a partialResult. */
-    instance.partialResult = partialResult;
-  }
+  Object.assign(instance, attributes ?? {});
   return instance;
 }
diff --git a/src/data-api/types/insert/insert-many.ts b/src/data-api/types/insert/insert-many.ts
@@ -12,15 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import type { IdOf } from '@/src/data-api/types';
+import { IdOf, SomeDoc } from '@/src/data-api/types';
 import { WithTimeout } from '@/src/common/types';
-import { SomeDoc } from '@/src/data-api';
+import { DataAPIErrorDescriptor } from '@/src/data-api';
 
 /** @internal */
 export interface InsertManyCommand {
   insertMany: {
     documents: unknown[],
     options?: {
+      returnDocumentResponses: true,
       ordered?: boolean,
     },
   }
@@ -176,3 +177,27 @@ export interface InsertManyResult<Schema extends SomeDoc> {
    */
   insertedCount: number;
 }
+
+/**
+ * Represents the specific status and id for a document present in the `insertMany` command
+ */
+export interface InsertManyDocumentResponse<Schema extends SomeDoc> {
+  /**
+   * The exact value of the `_id` field of the document that was inserted, whether it be the value passed by the client,
+   * or a server generated ID.
+   */
+  _id: IdOf<Schema>,
+  /**
+   * The processing status of the document
+   * - `OK`: The document was successfully processed, in which case the `error` field will be undefined for this object
+   * - `ERROR`: There was an error processing the document, in which case the `error` field will be present for this object
+   * - `SKIPPED`: The document was not processed because either the `insertMany` command was processing documents in order
+   * which means the processing fails at the first failure, or some other failure occurred before this document was
+   * processed. The `error` field will be undefined for this object.
+   */
+  status: 'OK' | 'ERROR' | 'SKIPPED',
+  /**
+   * The error which caused this document to fail insertion.
+   */
+  error?: DataAPIErrorDescriptor,
+}
diff --git a/tests/integration/data-api/collection/insert-many.test.ts b/tests/integration/data-api/collection/insert-many.test.ts
@@ -173,6 +173,7 @@ parallel('integration.data-api.collection.insert-many', { truncateColls: 'defaul
     assert.ok(error instanceof InsertManyError);
     assert.strictEqual(error.errorDescriptors[0].errorCode, 'DOCUMENT_ALREADY_EXISTS');
     assert.strictEqual(error.partialResult.insertedCount, 10);
+    assert.strictEqual(error.failedCount, 10);
     docs.slice(0, 10).forEach((doc, index) => {
       assert.strictEqual(error.partialResult.insertedIds[index], doc._id);
     });
@@ -191,6 +192,7 @@ parallel('integration.data-api.collection.insert-many', { truncateColls: 'defaul
     assert.ok(error instanceof InsertManyError);
     assert.strictEqual(error.errorDescriptors[0].errorCode, 'DOCUMENT_ALREADY_EXISTS');
     assert.strictEqual(error.partialResult.insertedCount, 19);
+    assert.strictEqual(error.failedCount, 1);
     docs.slice(0, 9).concat(docs.slice(10)).forEach((doc) => {
       assert.ok(error.partialResult.insertedIds.includes(doc._id));
     });
