opusdev/vector-similarity-api

1import { type Connection } from '../../cmap/connection';2import { ClientBulkWriteCursorResponse } from '../../cmap/wire_protocol/responses';3import type { ClientSession } from '../../sessions';4import { MongoDBNamespace } from '../../utils';5import { CommandOperation } from '../command';6import { Aspect, defineAspects } from '../operation';7import { type ClientBulkWriteCommand, type ClientBulkWriteCommandBuilder } from './command_builder';8import { type ClientBulkWriteOptions } from './common';9 10/**11 * Executes a single client bulk write operation within a potential batch.12 * @internal13 */14export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCursorResponse> {15  override SERVER_COMMAND_RESPONSE_TYPE = ClientBulkWriteCursorResponse;16 17  commandBuilder: ClientBulkWriteCommandBuilder;18  override options: ClientBulkWriteOptions;19 20  override get commandName() {21    return 'bulkWrite' as const;22  }23 24  constructor(commandBuilder: ClientBulkWriteCommandBuilder, options: ClientBulkWriteOptions) {25    super(undefined, options);26    this.commandBuilder = commandBuilder;27    this.options = options;28    this.ns = new MongoDBNamespace('admin', '$cmd');29  }30 31  override resetBatch(): boolean {32    return this.commandBuilder.resetBatch();33  }34 35  override get canRetryWrite(): boolean {36    return this.commandBuilder.isBatchRetryable;37  }38 39  override handleOk(40    response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>41  ): ClientBulkWriteCursorResponse {42    return response;43  }44 45  override buildCommandDocument(46    connection: Connection,47    _session?: ClientSession48  ): ClientBulkWriteCommand {49    const command = this.commandBuilder.buildBatch(50      connection.description.maxMessageSizeBytes,51      connection.description.maxWriteBatchSize,52      connection.description.maxBsonObjectSize53    );54 55    // Check _after_ the batch is built if we cannot retry it and override the option.56    if (!this.canRetryWrite) {57      this.options.willRetryWrite = false;58    }59 60    return command;61  }62}63 64// Skipping the collation as it goes on the individual ops.65defineAspects(ClientBulkWriteOperation, [66  Aspect.WRITE_OPERATION,67  Aspect.SKIP_COLLATION,68  Aspect.CURSOR_CREATING,69  Aspect.RETRYABLE,70  Aspect.COMMAND_BATCHING,71  Aspect.SUPPORTS_RAW_DATA72]);73