opusdev/vector-similarity-api

1import { BSON, type BSONSerializeOptions, type Document } from '../../bson';2import { DocumentSequence } from '../../cmap/commands';3import { MongoAPIError, MongoInvalidArgumentError } from '../../error';4import { type PkFactory } from '../../mongo_client';5import type { Filter, OptionalId, UpdateFilter, WithoutId } from '../../mongo_types';6import { formatSort, type SortForCmd } from '../../sort';7import { DEFAULT_PK_FACTORY, hasAtomicOperators } from '../../utils';8import { type CollationOptions } from '../command';9import { type Hint } from '../operation';10import type {11  AnyClientBulkWriteModel,12  ClientBulkWriteOptions,13  ClientDeleteManyModel,14  ClientDeleteOneModel,15  ClientInsertOneModel,16  ClientReplaceOneModel,17  ClientUpdateManyModel,18  ClientUpdateOneModel19} from './common';20 21/** @internal */22export interface ClientBulkWriteCommand {23  bulkWrite: 1;24  errorsOnly: boolean;25  ordered: boolean;26  ops: DocumentSequence;27  nsInfo: DocumentSequence;28  bypassDocumentValidation?: boolean;29  let?: Document;30  comment?: any;31}32 33/**34 * The bytes overhead for the extra fields added post command generation.35 */36const MESSAGE_OVERHEAD_BYTES = 1000;37 38/** @internal */39export class ClientBulkWriteCommandBuilder {40  models: ReadonlyArray<AnyClientBulkWriteModel<Document>>;41  options: ClientBulkWriteOptions;42  pkFactory: PkFactory;43  /** The current index in the models array that is being processed. */44  currentModelIndex: number;45  /** The model index that the builder was on when it finished the previous batch. Used for resets when retrying. */46  previousModelIndex: number;47  /** The last array of operations that were created. Used by the results merger for indexing results. */48  lastOperations: Document[];49  /** Returns true if the current batch being created has no multi-updates. */50  isBatchRetryable: boolean;51 52  /**53   * Create the command builder.54   * @param models - The client write models.55   */56  constructor(57    models: ReadonlyArray<AnyClientBulkWriteModel<Document>>,58    options: ClientBulkWriteOptions,59    pkFactory?: PkFactory60  ) {61    this.models = models;62    this.options = options;63    this.pkFactory = pkFactory ?? DEFAULT_PK_FACTORY;64    this.currentModelIndex = 0;65    this.previousModelIndex = 0;66    this.lastOperations = [];67    this.isBatchRetryable = true;68  }69 70  /**71   * Gets the errorsOnly value for the command, which is the inverse of the72   * user provided verboseResults option. Defaults to true.73   */74  get errorsOnly(): boolean {75    if ('verboseResults' in this.options) {76      return !this.options.verboseResults;77    }78    return true;79  }80 81  /**82   * Determines if there is another batch to process.83   * @returns True if not all batches have been built.84   */85  hasNextBatch(): boolean {86    return this.currentModelIndex < this.models.length;87  }88 89  /**90   * When we need to retry a command we need to set the current91   * model index back to its previous value.92   */93  resetBatch(): boolean {94    this.currentModelIndex = this.previousModelIndex;95    return true;96  }97 98  /**99   * Build a single batch of a client bulk write command.100   * @param maxMessageSizeBytes - The max message size in bytes.101   * @param maxWriteBatchSize - The max write batch size.102   * @returns The client bulk write command.103   */104  buildBatch(105    maxMessageSizeBytes: number,106    maxWriteBatchSize: number,107    maxBsonObjectSize: number108  ): ClientBulkWriteCommand {109    // We start by assuming the batch has no multi-updates, so it is retryable110    // until we find them.111    this.isBatchRetryable = true;112    let commandLength = 0;113    let currentNamespaceIndex = 0;114    const command: ClientBulkWriteCommand = this.baseCommand();115    const namespaces = new Map<string, number>();116    // In the case of retries we need to mark where we started this batch.117    this.previousModelIndex = this.currentModelIndex;118 119    while (this.currentModelIndex < this.models.length) {120      const model = this.models[this.currentModelIndex];121      const ns = model.namespace;122      const nsIndex = namespaces.get(ns);123 124      // Multi updates are not retryable.125      if (model.name === 'deleteMany' || model.name === 'updateMany') {126        this.isBatchRetryable = false;127      }128 129      if (nsIndex != null) {130        // Build the operation and serialize it to get the bytes buffer.131        const operation = buildOperation(model, nsIndex, this.pkFactory, this.options);132        let operationBuffer;133        try {134          operationBuffer = BSON.serialize(operation);135        } catch (cause) {136          throw new MongoInvalidArgumentError(`Could not serialize operation to BSON`, { cause });137        }138 139        validateBufferSize('ops', operationBuffer, maxBsonObjectSize);140 141        // Check if the operation buffer can fit in the command. If it can,142        // then add the operation to the document sequence and increment the143        // current length as long as the ops don't exceed the maxWriteBatchSize.144        if (145          commandLength + operationBuffer.length < maxMessageSizeBytes &&146          command.ops.documents.length < maxWriteBatchSize147        ) {148          // Pushing to the ops document sequence returns the total byte length of the document sequence.149          commandLength = MESSAGE_OVERHEAD_BYTES + command.ops.push(operation, operationBuffer);150          // Increment the builder's current model index.151          this.currentModelIndex++;152        } else {153          // The operation cannot fit in the current command and will need to154          // go in the next batch. Exit the loop.155          break;156        }157      } else {158        // The namespace is not already in the nsInfo so we will set it in the map, and159        // construct our nsInfo and ops documents and buffers.160        namespaces.set(ns, currentNamespaceIndex);161        const nsInfo = { ns: ns };162        const operation = buildOperation(163          model,164          currentNamespaceIndex,165          this.pkFactory,166          this.options167        );168        let nsInfoBuffer;169        let operationBuffer;170        try {171          nsInfoBuffer = BSON.serialize(nsInfo);172          operationBuffer = BSON.serialize(operation);173        } catch (cause) {174          throw new MongoInvalidArgumentError(`Could not serialize ns info to BSON`, { cause });175        }176 177        validateBufferSize('nsInfo', nsInfoBuffer, maxBsonObjectSize);178        validateBufferSize('ops', operationBuffer, maxBsonObjectSize);179 180        // Check if the operation and nsInfo buffers can fit in the command. If they181        // can, then add the operation and nsInfo to their respective document182        // sequences and increment the current length as long as the ops don't exceed183        // the maxWriteBatchSize.184        if (185          commandLength + nsInfoBuffer.length + operationBuffer.length < maxMessageSizeBytes &&186          command.ops.documents.length < maxWriteBatchSize187        ) {188          // Pushing to the ops document sequence returns the total byte length of the document sequence.189          commandLength =190            MESSAGE_OVERHEAD_BYTES +191            command.nsInfo.push(nsInfo, nsInfoBuffer) +192            command.ops.push(operation, operationBuffer);193          // We've added a new namespace, increment the namespace index.194          currentNamespaceIndex++;195          // Increment the builder's current model index.196          this.currentModelIndex++;197        } else {198          // The operation cannot fit in the current command and will need to199          // go in the next batch. Exit the loop.200          break;201        }202      }203    }204    // Set the last operations and return the command.205    this.lastOperations = command.ops.documents;206    return command;207  }208 209  private baseCommand(): ClientBulkWriteCommand {210    const command: ClientBulkWriteCommand = {211      bulkWrite: 1,212      errorsOnly: this.errorsOnly,213      ordered: this.options.ordered ?? true,214      ops: new DocumentSequence('ops'),215      nsInfo: new DocumentSequence('nsInfo')216    };217    // Add bypassDocumentValidation if it was present in the options.218    if (this.options.bypassDocumentValidation != null) {219      command.bypassDocumentValidation = this.options.bypassDocumentValidation;220    }221    // Add let if it was present in the options.222    if (this.options.let) {223      command.let = this.options.let;224    }225 226    // we check for undefined specifically here to allow falsy values227    // eslint-disable-next-line no-restricted-syntax228    if (this.options.comment !== undefined) {229      command.comment = this.options.comment;230    }231 232    return command;233  }234}235 236function validateBufferSize(name: string, buffer: Uint8Array, maxBsonObjectSize: number) {237  if (buffer.length > maxBsonObjectSize) {238    throw new MongoInvalidArgumentError(239      `Client bulk write operation ${name} of length ${buffer.length} exceeds the max bson object size of ${maxBsonObjectSize}`240    );241  }242}243 244/** @internal */245interface ClientInsertOperation {246  insert: number;247  document: OptionalId<Document>;248}249 250/**251 * Build the insert one operation.252 * @param model - The insert one model.253 * @param index - The namespace index.254 * @returns the operation.255 */256export const buildInsertOneOperation = (257  model: ClientInsertOneModel<Document>,258  index: number,259  pkFactory: PkFactory260): ClientInsertOperation => {261  const document: ClientInsertOperation = {262    insert: index,263    document: model.document264  };265  document.document._id = model.document._id ?? pkFactory.createPk();266  return document;267};268 269/** @internal */270export interface ClientDeleteOperation {271  delete: number;272  multi: boolean;273  filter: Filter<Document>;274  hint?: Hint;275  collation?: CollationOptions;276}277 278/**279 * Build the delete one operation.280 * @param model - The insert many model.281 * @param index - The namespace index.282 * @returns the operation.283 */284export const buildDeleteOneOperation = (285  model: ClientDeleteOneModel<Document>,286  index: number287): Document => {288  return createDeleteOperation(model, index, false);289};290 291/**292 * Build the delete many operation.293 * @param model - The delete many model.294 * @param index - The namespace index.295 * @returns the operation.296 */297export const buildDeleteManyOperation = (298  model: ClientDeleteManyModel<Document>,299  index: number300): Document => {301  return createDeleteOperation(model, index, true);302};303 304/**305 * Creates a delete operation based on the parameters.306 */307function createDeleteOperation(308  model: ClientDeleteOneModel<Document> | ClientDeleteManyModel<Document>,309  index: number,310  multi: boolean311): ClientDeleteOperation {312  const document: ClientDeleteOperation = {313    delete: index,314    multi: multi,315    filter: model.filter316  };317  if (model.hint) {318    document.hint = model.hint;319  }320  if (model.collation) {321    document.collation = model.collation;322  }323  return document;324}325 326/** @internal */327export interface ClientUpdateOperation {328  update: number;329  multi: boolean;330  filter: Filter<Document>;331  updateMods: UpdateFilter<Document> | Document[];332  hint?: Hint;333  upsert?: boolean;334  arrayFilters?: Document[];335  collation?: CollationOptions;336  sort?: SortForCmd;337}338 339/**340 * Build the update one operation.341 * @param model - The update one model.342 * @param index - The namespace index.343 * @returns the operation.344 */345export const buildUpdateOneOperation = (346  model: ClientUpdateOneModel<Document>,347  index: number,348  options: BSONSerializeOptions349): ClientUpdateOperation => {350  return createUpdateOperation(model, index, false, options);351};352 353/**354 * Build the update many operation.355 * @param model - The update many model.356 * @param index - The namespace index.357 * @returns the operation.358 */359export const buildUpdateManyOperation = (360  model: ClientUpdateManyModel<Document>,361  index: number,362  options: BSONSerializeOptions363): ClientUpdateOperation => {364  return createUpdateOperation(model, index, true, options);365};366 367/**368 * Validate the update document.369 * @param update - The update document.370 */371function validateUpdate(update: Document, options: BSONSerializeOptions) {372  if (!hasAtomicOperators(update, options)) {373    throw new MongoAPIError(374      'Client bulk write update models must only contain atomic modifiers (start with $) and must not be empty.'375    );376  }377}378 379/**380 * Creates a delete operation based on the parameters.381 */382function createUpdateOperation(383  model: ClientUpdateOneModel<Document> | ClientUpdateManyModel<Document>,384  index: number,385  multi: boolean,386  options: BSONSerializeOptions387): ClientUpdateOperation {388  // Update documents provided in UpdateOne and UpdateMany write models are389  // required only to contain atomic modifiers (i.e. keys that start with "$").390  // Drivers MUST throw an error if an update document is empty or if the391  // document's first key does not start with "$".392  validateUpdate(model.update, options);393  const document: ClientUpdateOperation = {394    update: index,395    multi: multi,396    filter: model.filter,397    updateMods: model.update398  };399  if (model.hint) {400    document.hint = model.hint;401  }402  if (model.upsert) {403    document.upsert = model.upsert;404  }405  if (model.arrayFilters) {406    document.arrayFilters = model.arrayFilters;407  }408  if (model.collation) {409    document.collation = model.collation;410  }411  if (!multi && 'sort' in model && model.sort != null) {412    document.sort = formatSort(model.sort);413  }414  return document;415}416 417/** @internal */418export interface ClientReplaceOneOperation {419  update: number;420  multi: boolean;421  filter: Filter<Document>;422  updateMods: WithoutId<Document>;423  hint?: Hint;424  upsert?: boolean;425  collation?: CollationOptions;426  sort?: SortForCmd;427}428 429/**430 * Build the replace one operation.431 * @param model - The replace one model.432 * @param index - The namespace index.433 * @returns the operation.434 */435export const buildReplaceOneOperation = (436  model: ClientReplaceOneModel<Document>,437  index: number438): ClientReplaceOneOperation => {439  if (hasAtomicOperators(model.replacement)) {440    throw new MongoAPIError(441      'Client bulk write replace models must not contain atomic modifiers (start with $) and must not be empty.'442    );443  }444 445  const document: ClientReplaceOneOperation = {446    update: index,447    multi: false,448    filter: model.filter,449    updateMods: model.replacement450  };451  if (model.hint) {452    document.hint = model.hint;453  }454  if (model.upsert) {455    document.upsert = model.upsert;456  }457  if (model.collation) {458    document.collation = model.collation;459  }460  if (model.sort != null) {461    document.sort = formatSort(model.sort);462  }463  return document;464};465 466/** @internal */467export function buildOperation(468  model: AnyClientBulkWriteModel<Document>,469  index: number,470  pkFactory: PkFactory,471  options: BSONSerializeOptions472): Document {473  switch (model.name) {474    case 'insertOne':475      return buildInsertOneOperation(model, index, pkFactory);476    case 'deleteOne':477      return buildDeleteOneOperation(model, index);478    case 'deleteMany':479      return buildDeleteManyOperation(model, index);480    case 'updateOne':481      return buildUpdateOneOperation(model, index, options);482    case 'updateMany':483      return buildUpdateManyOperation(model, index, options);484    case 'replaceOne':485      return buildReplaceOneOperation(model, index);486  }487}488