opusdev/vector-similarity-api

1import { MongoWriteConcernError } from '../..';2import { type Document } from '../../bson';3import { type ClientBulkWriteCursor } from '../../cursor/client_bulk_write_cursor';4import { MongoClientBulkWriteError } from '../../error';5import {6  type ClientBulkWriteError,7  type ClientBulkWriteOptions,8  type ClientBulkWriteResult,9  type ClientDeleteResult,10  type ClientInsertOneResult,11  type ClientUpdateResult12} from './common';13 14/**15 * Unacknowledged bulk writes are always the same.16 */17const UNACKNOWLEDGED = {18  acknowledged: false,19  insertedCount: 0,20  upsertedCount: 0,21  matchedCount: 0,22  modifiedCount: 0,23  deletedCount: 0,24  insertResults: undefined,25  updateResults: undefined,26  deleteResults: undefined27};28 29interface ClientBulkWriteResultAccumulation {30  /**31   * Whether the bulk write was acknowledged.32   */33  acknowledged: boolean;34  /**35   * The total number of documents inserted across all insert operations.36   */37  insertedCount: number;38  /**39   * The total number of documents upserted across all update operations.40   */41  upsertedCount: number;42  /**43   * The total number of documents matched across all update operations.44   */45  matchedCount: number;46  /**47   * The total number of documents modified across all update operations.48   */49  modifiedCount: number;50  /**51   * The total number of documents deleted across all delete operations.52   */53  deletedCount: number;54  /**55   * The results of each individual insert operation that was successfully performed.56   */57  insertResults?: Map<number, ClientInsertOneResult>;58  /**59   * The results of each individual update operation that was successfully performed.60   */61  updateResults?: Map<number, ClientUpdateResult>;62  /**63   * The results of each individual delete operation that was successfully performed.64   */65  deleteResults?: Map<number, ClientDeleteResult>;66}67 68/**69 * Merges client bulk write cursor responses together into a single result.70 * @internal71 */72export class ClientBulkWriteResultsMerger {73  private result: ClientBulkWriteResultAccumulation;74  private options: ClientBulkWriteOptions;75  private currentBatchOffset: number;76  writeConcernErrors: Document[];77  writeErrors: Map<number, ClientBulkWriteError>;78 79  /**80   * @returns The standard unacknowledged bulk write result.81   */82  static unacknowledged(): ClientBulkWriteResult {83    return UNACKNOWLEDGED;84  }85 86  /**87   * Instantiate the merger.88   * @param options - The options.89   */90  constructor(options: ClientBulkWriteOptions) {91    this.options = options;92    this.currentBatchOffset = 0;93    this.writeConcernErrors = [];94    this.writeErrors = new Map();95    this.result = {96      acknowledged: true,97      insertedCount: 0,98      upsertedCount: 0,99      matchedCount: 0,100      modifiedCount: 0,101      deletedCount: 0,102      insertResults: undefined,103      updateResults: undefined,104      deleteResults: undefined105    };106 107    if (options.verboseResults) {108      this.result.insertResults = new Map<number, ClientInsertOneResult>();109      this.result.updateResults = new Map<number, ClientUpdateResult>();110      this.result.deleteResults = new Map<number, ClientDeleteResult>();111    }112  }113 114  /**115   * Get the bulk write result object.116   */117  get bulkWriteResult(): ClientBulkWriteResult {118    return {119      acknowledged: this.result.acknowledged,120      insertedCount: this.result.insertedCount,121      upsertedCount: this.result.upsertedCount,122      matchedCount: this.result.matchedCount,123      modifiedCount: this.result.modifiedCount,124      deletedCount: this.result.deletedCount,125      insertResults: this.result.insertResults,126      updateResults: this.result.updateResults,127      deleteResults: this.result.deleteResults128    };129  }130 131  /**132   * Merge the results in the cursor to the existing result.133   * @param currentBatchOffset - The offset index to the original models.134   * @param response - The cursor response.135   * @param documents - The documents in the cursor.136   * @returns The current result.137   */138  async merge(cursor: ClientBulkWriteCursor): Promise<ClientBulkWriteResult> {139    let writeConcernErrorResult;140    try {141      for await (const document of cursor) {142        // Only add to maps if ok: 1143        if (document.ok === 1) {144          if (this.options.verboseResults) {145            this.processDocument(cursor, document);146          }147        } else {148          // If an individual write error is encountered during an ordered bulk write, drivers MUST149          // record the error in writeErrors and immediately throw the exception. Otherwise, drivers150          // MUST continue to iterate the results cursor and execute any further bulkWrite batches.151          if (this.options.ordered) {152            const error = new MongoClientBulkWriteError({153              message: 'Mongo client ordered bulk write encountered a write error.'154            });155            error.writeErrors.set(document.idx + this.currentBatchOffset, {156              code: document.code,157              message: document.errmsg158            });159            error.partialResult = this.result;160            throw error;161          } else {162            this.writeErrors.set(document.idx + this.currentBatchOffset, {163              code: document.code,164              message: document.errmsg165            });166          }167        }168      }169    } catch (error) {170      if (error instanceof MongoWriteConcernError) {171        const result = error.result;172        writeConcernErrorResult = {173          insertedCount: result.nInserted,174          upsertedCount: result.nUpserted,175          matchedCount: result.nMatched,176          modifiedCount: result.nModified,177          deletedCount: result.nDeleted,178          writeConcernError: result.writeConcernError179        };180        if (this.options.verboseResults && result.cursor.firstBatch) {181          for (const document of result.cursor.firstBatch) {182            if (document.ok === 1) {183              this.processDocument(cursor, document);184            }185          }186        }187      } else {188        throw error;189      }190    } finally {191      // Update the counts from the cursor response.192      if (cursor.response) {193        const response = cursor.response;194        this.incrementCounts(response);195      }196 197      // Increment the batch offset.198      this.currentBatchOffset += cursor.operations.length;199    }200 201    // If we have write concern errors ensure they are added.202    if (writeConcernErrorResult) {203      const writeConcernError = writeConcernErrorResult.writeConcernError as Document;204      this.incrementCounts(writeConcernErrorResult);205      this.writeConcernErrors.push({206        code: writeConcernError.code,207        message: writeConcernError.errmsg208      });209    }210 211    return this.result;212  }213 214  /**215   * Process an individual document in the results.216   * @param cursor - The cursor.217   * @param document - The document to process.218   */219  private processDocument(cursor: ClientBulkWriteCursor, document: Document) {220    // Get the corresponding operation from the command.221    const operation = cursor.operations[document.idx];222    // Handle insert results.223    if ('insert' in operation) {224      this.result.insertResults?.set(document.idx + this.currentBatchOffset, {225        insertedId: operation.document._id226      });227    }228    // Handle update results.229    if ('update' in operation) {230      const result: ClientUpdateResult = {231        matchedCount: document.n,232        modifiedCount: document.nModified ?? 0,233        // Check if the bulk did actually upsert.234        didUpsert: document.upserted != null235      };236      if (document.upserted) {237        result.upsertedId = document.upserted._id;238      }239      this.result.updateResults?.set(document.idx + this.currentBatchOffset, result);240    }241    // Handle delete results.242    if ('delete' in operation) {243      this.result.deleteResults?.set(document.idx + this.currentBatchOffset, {244        deletedCount: document.n245      });246    }247  }248 249  /**250   * Increment the result counts.251   * @param document - The document with the results.252   */253  private incrementCounts(document: Document) {254    this.result.insertedCount += document.insertedCount;255    this.result.upsertedCount += document.upsertedCount;256    this.result.matchedCount += document.matchedCount;257    this.result.modifiedCount += document.modifiedCount;258    this.result.deletedCount += document.deletedCount;259  }260}261