opusdev/vector-similarity-api

1import type { Document } from '../bson';2import * as BSON from '../bson';3import type { Collection } from '../collection';4import { MongoInvalidArgumentError } from '../error';5import type { DeleteStatement } from '../operations/delete';6import type { UpdateStatement } from '../operations/update';7import {8  Batch,9  BatchType,10  BulkOperationBase,11  type BulkWriteOptions,12  type BulkWriteResult13} from './common';14 15/** @public */16export class UnorderedBulkOperation extends BulkOperationBase {17  /** @internal */18  constructor(collection: Collection, options: BulkWriteOptions) {19    super(collection, options, false);20  }21 22  override handleWriteError(writeResult: BulkWriteResult): void {23    if (this.s.batches.length) {24      return;25    }26 27    return super.handleWriteError(writeResult);28  }29 30  addToOperationsList(31    batchType: BatchType,32    document: Document | UpdateStatement | DeleteStatement33  ): this {34    // Get the bsonSize35    const bsonSize = BSON.calculateObjectSize(document, {36      checkKeys: false,37 38      // Since we don't know what the user selected for BSON options here,39      // err on the safe side, and check the size with ignoreUndefined: false.40      ignoreUndefined: false41    } as any);42 43    // Throw error if the doc is bigger than the max BSON size44    if (bsonSize >= this.s.maxBsonObjectSize) {45      // TODO(NODE-3483): Change this to MongoBSONError46      throw new MongoInvalidArgumentError(47        `Document is larger than the maximum size ${this.s.maxBsonObjectSize}`48      );49    }50 51    // Holds the current batch52    this.s.currentBatch = undefined;53    // Get the right type of batch54    if (batchType === BatchType.INSERT) {55      this.s.currentBatch = this.s.currentInsertBatch;56    } else if (batchType === BatchType.UPDATE) {57      this.s.currentBatch = this.s.currentUpdateBatch;58    } else if (batchType === BatchType.DELETE) {59      this.s.currentBatch = this.s.currentRemoveBatch;60    }61 62    const maxKeySize = this.s.maxKeySize;63 64    // Create a new batch object if we don't have a current one65    if (this.s.currentBatch == null) {66      this.s.currentBatch = new Batch(batchType, this.s.currentIndex);67    }68 69    // Check if we need to create a new batch70    if (71      // New batch if we exceed the max batch op size72      this.s.currentBatch.size + 1 >= this.s.maxWriteBatchSize ||73      // New batch if we exceed the maxBatchSizeBytes. Only matters if batch already has a doc,74      // since we can't sent an empty batch75      (this.s.currentBatch.size > 0 &&76        this.s.currentBatch.sizeBytes + maxKeySize + bsonSize >= this.s.maxBatchSizeBytes) ||77      // New batch if the new op does not have the same op type as the current batch78      this.s.currentBatch.batchType !== batchType79    ) {80      // Save the batch to the execution stack81      this.s.batches.push(this.s.currentBatch);82 83      // Create a new batch84      this.s.currentBatch = new Batch(batchType, this.s.currentIndex);85    }86 87    // We have an array of documents88    if (Array.isArray(document)) {89      throw new MongoInvalidArgumentError('Operation passed in cannot be an Array');90    }91 92    this.s.currentBatch.operations.push(document);93    this.s.currentBatch.originalIndexes.push(this.s.currentIndex);94    this.s.currentIndex = this.s.currentIndex + 1;95 96    // Save back the current Batch to the right type97    if (batchType === BatchType.INSERT) {98      this.s.currentInsertBatch = this.s.currentBatch;99      this.s.bulkResult.insertedIds.push({100        index: this.s.bulkResult.insertedIds.length,101        _id: (document as Document)._id102      });103    } else if (batchType === BatchType.UPDATE) {104      this.s.currentUpdateBatch = this.s.currentBatch;105    } else if (batchType === BatchType.DELETE) {106      this.s.currentRemoveBatch = this.s.currentBatch;107    }108 109    // Update current batch size110    this.s.currentBatch.size += 1;111    this.s.currentBatch.sizeBytes += maxKeySize + bsonSize;112 113    return this;114  }115}116