opusdev/vector-similarity-api

1import { type Document } from '../../bson';2import { CursorTimeoutContext, CursorTimeoutMode } from '../../cursor/abstract_cursor';3import { ClientBulkWriteCursor } from '../../cursor/client_bulk_write_cursor';4import {5  MongoClientBulkWriteError,6  MongoClientBulkWriteExecutionError,7  MongoInvalidArgumentError,8  MongoServerError9} from '../../error';10import { type MongoClient } from '../../mongo_client';11import { TimeoutContext } from '../../timeout';12import { resolveTimeoutOptions } from '../../utils';13import { WriteConcern } from '../../write_concern';14import { executeOperation } from '../execute_operation';15import { ClientBulkWriteOperation } from './client_bulk_write';16import { ClientBulkWriteCommandBuilder } from './command_builder';17import {18  type AnyClientBulkWriteModel,19  type ClientBulkWriteOptions,20  type ClientBulkWriteResult21} from './common';22import { ClientBulkWriteResultsMerger } from './results_merger';23 24/**25 * Responsible for executing a client bulk write.26 * @internal27 */28export class ClientBulkWriteExecutor {29  private readonly client: MongoClient;30  private readonly options: ClientBulkWriteOptions;31  private readonly operations: ReadonlyArray<AnyClientBulkWriteModel<Document>>;32 33  /**34   * Instantiate the executor.35   * @param client - The mongo client.36   * @param operations - The user supplied bulk write models.37   * @param options - The bulk write options.38   */39  constructor(40    client: MongoClient,41    operations: ReadonlyArray<AnyClientBulkWriteModel<Document>>,42    options?: ClientBulkWriteOptions43  ) {44    if (operations.length === 0) {45      throw new MongoClientBulkWriteExecutionError('No client bulk write models were provided.');46    }47 48    this.client = client;49    this.operations = operations;50    this.options = {51      ordered: true,52      bypassDocumentValidation: false,53      verboseResults: false,54      ...options55    };56 57    // If no write concern was provided, we inherit one from the client.58    if (!this.options.writeConcern) {59      this.options.writeConcern = WriteConcern.fromOptions(this.client.s.options);60    }61 62    if (this.options.writeConcern?.w === 0) {63      if (this.options.verboseResults) {64        throw new MongoInvalidArgumentError(65          'Cannot request unacknowledged write concern and verbose results'66        );67      }68 69      if (this.options.ordered) {70        throw new MongoInvalidArgumentError(71          'Cannot request unacknowledged write concern and ordered writes'72        );73      }74    }75  }76 77  /**78   * Execute the client bulk write. Will split commands into batches and exhaust the cursors79   * for each, then merge the results into one.80   * @returns The result.81   */82  async execute(): Promise<ClientBulkWriteResult> {83    // The command builder will take the user provided models and potential split the batch84    // into multiple commands due to size.85    const pkFactory = this.client.s.options.pkFactory;86    const commandBuilder = new ClientBulkWriteCommandBuilder(87      this.operations,88      this.options,89      pkFactory90    );91    // Unacknowledged writes need to execute all batches and return { ok: 1}92    const resolvedOptions = resolveTimeoutOptions(this.client, this.options);93    const context = TimeoutContext.create(resolvedOptions);94 95    if (this.options.writeConcern?.w === 0) {96      while (commandBuilder.hasNextBatch()) {97        const operation = new ClientBulkWriteOperation(commandBuilder, this.options);98        await executeOperation(this.client, operation, context);99      }100      return ClientBulkWriteResultsMerger.unacknowledged();101    } else {102      const resultsMerger = new ClientBulkWriteResultsMerger(this.options);103      // For each command will will create and exhaust a cursor for the results.104      while (commandBuilder.hasNextBatch()) {105        const cursorContext = new CursorTimeoutContext(context, Symbol());106        const options = {107          ...this.options,108          timeoutContext: cursorContext,109          ...(resolvedOptions.timeoutMS != null && { timeoutMode: CursorTimeoutMode.LIFETIME })110        };111        const cursor = new ClientBulkWriteCursor(this.client, commandBuilder, options);112        try {113          await resultsMerger.merge(cursor);114        } catch (error) {115          // Write concern errors are recorded in the writeConcernErrors field on MongoClientBulkWriteError.116          // When a write concern error is encountered, it should not terminate execution of the bulk write117          // for either ordered or unordered bulk writes. However, drivers MUST throw an exception at the end118          // of execution if any write concern errors were observed.119          if (error instanceof MongoServerError && !(error instanceof MongoClientBulkWriteError)) {120            // Server side errors need to be wrapped inside a MongoClientBulkWriteError, where the root121            // cause is the error property and a partial result is to be included.122            const bulkWriteError = new MongoClientBulkWriteError({123              message: 'Mongo client bulk write encountered an error during execution'124            });125            bulkWriteError.cause = error;126            bulkWriteError.partialResult = resultsMerger.bulkWriteResult;127            throw bulkWriteError;128          } else {129            // Client side errors are just thrown.130            throw error;131          }132        }133      }134 135      // If we have write concern errors or unordered write errors at the end we throw.136      if (resultsMerger.writeConcernErrors.length > 0 || resultsMerger.writeErrors.size > 0) {137        const error = new MongoClientBulkWriteError({138          message: 'Mongo client bulk write encountered errors during execution.'139        });140        error.writeConcernErrors = resultsMerger.writeConcernErrors;141        error.writeErrors = resultsMerger.writeErrors;142        error.partialResult = resultsMerger.bulkWriteResult;143        throw error;144      }145 146      return resultsMerger.bulkWriteResult;147    }148  }149}150