mongodb 6.9.0 → 6.10.0

package/src/operations/client_bulk_write/client_bulk_write.ts

@@ -0,0 +1,107 @@
+import { MongoClientBulkWriteExecutionError, ServerType } from '../../beta';
+import { ClientBulkWriteCursorResponse } from '../../cmap/wire_protocol/responses';
+import type { Server } from '../../sdam/server';
+import type { ClientSession } from '../../sessions';
+import { MongoDBNamespace } from '../../utils';
+import { CommandOperation } from '../command';
+import { Aspect, defineAspects } from '../operation';
+import { type ClientBulkWriteCommandBuilder } from './command_builder';
+import { type ClientBulkWriteOptions } from './common';
+
+/**
+ * Executes a single client bulk write operation within a potential batch.
+ * @internal
+ */
+export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCursorResponse> {
+  commandBuilder: ClientBulkWriteCommandBuilder;
+  override options: ClientBulkWriteOptions;
+
+  override get commandName() {
+    return 'bulkWrite' as const;
+  }
+
+  constructor(commandBuilder: ClientBulkWriteCommandBuilder, options: ClientBulkWriteOptions) {
+    super(undefined, options);
+    this.commandBuilder = commandBuilder;
+    this.options = options;
+    this.ns = new MongoDBNamespace('admin', '$cmd');
+  }
+
+  override resetBatch(): boolean {
+    return this.commandBuilder.resetBatch();
+  }
+
+  override get canRetryWrite(): boolean {
+    return this.commandBuilder.isBatchRetryable;
+  }
+
+  /**
+   * Execute the command. Superclass will handle write concern, etc.
+   * @param server - The server.
+   * @param session - The session.
+   * @returns The response.
+   */
+  override async execute(
+    server: Server,
+    session: ClientSession | undefined
+  ): Promise<ClientBulkWriteCursorResponse> {
+    let command;
+
+    if (server.description.type === ServerType.LoadBalancer) {
+      if (session) {
+        let connection;
+        if (!session.pinnedConnection) {
+          // Checkout a connection to build the command.
+          connection = await server.pool.checkOut();
+          // Pin the connection to the session so it get used to execute the command and we do not
+          // perform a double check-in/check-out.
+          session.pin(connection);
+        } else {
+          connection = session.pinnedConnection;
+        }
+        command = this.commandBuilder.buildBatch(
+          connection.hello?.maxMessageSizeBytes,
+          connection.hello?.maxWriteBatchSize,
+          connection.hello?.maxBsonObjectSize
+        );
+      } else {
+        throw new MongoClientBulkWriteExecutionError(
+          'Session provided to the client bulk write operation must be present.'
+        );
+      }
+    } else {
+      // At this point we have a server and the auto connect code has already
+      // run in executeOperation, so the server description will be populated.
+      // We can use that to build the command.
+      if (
+        !server.description.maxWriteBatchSize ||
+        !server.description.maxMessageSizeBytes ||
+        !server.description.maxBsonObjectSize
+      ) {
+        throw new MongoClientBulkWriteExecutionError(
+          'In order to execute a client bulk write, both maxWriteBatchSize, maxMessageSizeBytes and maxBsonObjectSize must be provided by the servers hello response.'
+        );
+      }
+      command = this.commandBuilder.buildBatch(
+        server.description.maxMessageSizeBytes,
+        server.description.maxWriteBatchSize,
+        server.description.maxBsonObjectSize
+      );
+    }
+
+    // Check after the batch is built if we cannot retry it and override the option.
+    if (!this.canRetryWrite) {
+      this.options.willRetryWrite = false;
+    }
+    return await super.executeCommand(server, session, command, ClientBulkWriteCursorResponse);
+  }
+}
+
+// Skipping the collation as it goes on the individual ops.
+defineAspects(ClientBulkWriteOperation, [
+  Aspect.WRITE_OPERATION,
+  Aspect.SKIP_COLLATION,
+  Aspect.CURSOR_CREATING,
+  Aspect.RETRYABLE,
+  Aspect.COMMAND_BATCHING
+]);

package/src/operations/client_bulk_write/command_builder.ts

@@ -1,6 +1,9 @@
-import { type Document } from '../../bson';
+import { BSON, type Document } from '../../bson';
 import { DocumentSequence } from '../../cmap/commands';
+import { MongoAPIError, MongoInvalidArgumentError } from '../../error';
+import { type PkFactory } from '../../mongo_client';
 import type { Filter, OptionalId, UpdateFilter, WithoutId } from '../../mongo_types';
+import { DEFAULT_PK_FACTORY, hasAtomicOperators } from '../../utils';
 import { type CollationOptions } from '../command';
 import { type Hint } from '../operation';
 import type {
@@ -23,20 +26,44 @@ export interface ClientBulkWriteCommand {
   nsInfo: DocumentSequence;
   bypassDocumentValidation?: boolean;
   let?: Document;
+  comment?: any;
 }
 
+/**
+ * The bytes overhead for the extra fields added post command generation.
+ */
+const MESSAGE_OVERHEAD_BYTES = 1000;
+
 /** @internal */
 export class ClientBulkWriteCommandBuilder {
-  models: AnyClientBulkWriteModel[];
+  models: ReadonlyArray<AnyClientBulkWriteModel<Document>>;
   options: ClientBulkWriteOptions;
+  pkFactory: PkFactory;
+  /** The current index in the models array that is being processed. */
+  currentModelIndex: number;
+  /** The model index that the builder was on when it finished the previous batch. Used for resets when retrying. */
+  previousModelIndex: number;
+  /** The last array of operations that were created. Used by the results merger for indexing results. */
+  lastOperations: Document[];
+  /** Returns true if the current batch being created has no multi-updates. */
+  isBatchRetryable: boolean;
 
   /**
    * Create the command builder.
    * @param models - The client write models.
    */
-  constructor(models: AnyClientBulkWriteModel[], options: ClientBulkWriteOptions) {
+  constructor(
+    models: ReadonlyArray<AnyClientBulkWriteModel<Document>>,
+    options: ClientBulkWriteOptions,
+    pkFactory?: PkFactory
+  ) {
     this.models = models;
     this.options = options;
+    this.pkFactory = pkFactory ?? DEFAULT_PK_FACTORY;
+    this.currentModelIndex = 0;
+    this.previousModelIndex = 0;
+    this.lastOperations = [];
+    this.isBatchRetryable = true;
   }
 
   /**
@@ -51,34 +78,135 @@ export class ClientBulkWriteCommandBuilder {
   }
 
   /**
-   * Build the bulk write commands from the models.
+   * Determines if there is another batch to process.
+   * @returns True if not all batches have been built.
+   */
+  hasNextBatch(): boolean {
+    return this.currentModelIndex < this.models.length;
+  }
+
+  /**
+   * When we need to retry a command we need to set the current
+   * model index back to its previous value.
+   */
+  resetBatch(): boolean {
+    this.currentModelIndex = this.previousModelIndex;
+    return true;
+  }
+
+  /**
+   * Build a single batch of a client bulk write command.
+   * @param maxMessageSizeBytes - The max message size in bytes.
+   * @param maxWriteBatchSize - The max write batch size.
+   * @returns The client bulk write command.
    */
-  buildCommands(): ClientBulkWriteCommand[] {
-    // Iterate the models to build the ops and nsInfo fields.
-    const operations = [];
+  buildBatch(
+    maxMessageSizeBytes: number,
+    maxWriteBatchSize: number,
+    maxBsonObjectSize: number
+  ): ClientBulkWriteCommand {
+    // We start by assuming the batch has no multi-updates, so it is retryable
+    // until we find them.
+    this.isBatchRetryable = true;
+    let commandLength = 0;
     let currentNamespaceIndex = 0;
+    const command: ClientBulkWriteCommand = this.baseCommand();
     const namespaces = new Map<string, number>();
-    for (const model of this.models) {
+    // In the case of retries we need to mark where we started this batch.
+    this.previousModelIndex = this.currentModelIndex;
+
+    while (this.currentModelIndex < this.models.length) {
+      const model = this.models[this.currentModelIndex];
       const ns = model.namespace;
-      const index = namespaces.get(ns);
-      if (index != null) {
-        operations.push(buildOperation(model, index));
+      const nsIndex = namespaces.get(ns);
+
+      // Multi updates are not retryable.
+      if (model.name === 'deleteMany' || model.name === 'updateMany') {
+        this.isBatchRetryable = false;
+      }
+
+      if (nsIndex != null) {
+        // Build the operation and serialize it to get the bytes buffer.
+        const operation = buildOperation(model, nsIndex, this.pkFactory);
+        let operationBuffer;
+        try {
+          operationBuffer = BSON.serialize(operation);
+        } catch (cause) {
+          throw new MongoInvalidArgumentError(`Could not serialize operation to BSON`, { cause });
+        }
+
+        validateBufferSize('ops', operationBuffer, maxBsonObjectSize);
+
+        // Check if the operation buffer can fit in the command. If it can,
+        // then add the operation to the document sequence and increment the
+        // current length as long as the ops don't exceed the maxWriteBatchSize.
+        if (
+          commandLength + operationBuffer.length < maxMessageSizeBytes &&
+          command.ops.documents.length < maxWriteBatchSize
+        ) {
+          // Pushing to the ops document sequence returns the total byte length of the document sequence.
+          commandLength = MESSAGE_OVERHEAD_BYTES + command.ops.push(operation, operationBuffer);
+          // Increment the builder's current model index.
+          this.currentModelIndex++;
+        } else {
+          // The operation cannot fit in the current command and will need to
+          // go in the next batch. Exit the loop.
+          break;
+        }
       } else {
+        // The namespace is not already in the nsInfo so we will set it in the map, and
+        // construct our nsInfo and ops documents and buffers.
         namespaces.set(ns, currentNamespaceIndex);
-        operations.push(buildOperation(model, currentNamespaceIndex));
-        currentNamespaceIndex++;
+        const nsInfo = { ns: ns };
+        const operation = buildOperation(model, currentNamespaceIndex, this.pkFactory);
+        let nsInfoBuffer;
+        let operationBuffer;
+        try {
+          nsInfoBuffer = BSON.serialize(nsInfo);
+          operationBuffer = BSON.serialize(operation);
+        } catch (cause) {
+          throw new MongoInvalidArgumentError(`Could not serialize ns info to BSON`, { cause });
+        }
+
+        validateBufferSize('nsInfo', nsInfoBuffer, maxBsonObjectSize);
+        validateBufferSize('ops', operationBuffer, maxBsonObjectSize);
+
+        // Check if the operation and nsInfo buffers can fit in the command. If they
+        // can, then add the operation and nsInfo to their respective document
+        // sequences and increment the current length as long as the ops don't exceed
+        // the maxWriteBatchSize.
+        if (
+          commandLength + nsInfoBuffer.length + operationBuffer.length < maxMessageSizeBytes &&
+          command.ops.documents.length < maxWriteBatchSize
+        ) {
+          // Pushing to the ops document sequence returns the total byte length of the document sequence.
+          commandLength =
+            MESSAGE_OVERHEAD_BYTES +
+            command.nsInfo.push(nsInfo, nsInfoBuffer) +
+            command.ops.push(operation, operationBuffer);
+          // We've added a new namespace, increment the namespace index.
+          currentNamespaceIndex++;
+          // Increment the builder's current model index.
+          this.currentModelIndex++;
+        } else {
+          // The operation cannot fit in the current command and will need to
+          // go in the next batch. Exit the loop.
+          break;
+        }
       }
     }
+    // Set the last operations and return the command.
+    this.lastOperations = command.ops.documents;
+    return command;
+  }
 
-    const nsInfo = Array.from(namespaces.keys(), ns => ({ ns }));
-
-    // The base command.
+  private baseCommand(): ClientBulkWriteCommand {
     const command: ClientBulkWriteCommand = {
       bulkWrite: 1,
       errorsOnly: this.errorsOnly,
       ordered: this.options.ordered ?? true,
-      ops: new DocumentSequence(operations),
-      nsInfo: new DocumentSequence(nsInfo)
+      ops: new DocumentSequence('ops'),
+      nsInfo: new DocumentSequence('nsInfo')
     };
     // Add bypassDocumentValidation if it was present in the options.
     if (this.options.bypassDocumentValidation != null) {
@@ -88,7 +216,22 @@ export class ClientBulkWriteCommandBuilder {
     if (this.options.let) {
       command.let = this.options.let;
     }
-    return [command];
+
+    // we check for undefined specifically here to allow falsy values
+    // eslint-disable-next-line no-restricted-syntax
+    if (this.options.comment !== undefined) {
+      command.comment = this.options.comment;
+    }
+
+    return command;
+  }
+}
+
+function validateBufferSize(name: string, buffer: Uint8Array, maxBsonObjectSize: number) {
+  if (buffer.length > maxBsonObjectSize) {
+    throw new MongoInvalidArgumentError(
+      `Client bulk write operation ${name} of length ${buffer.length} exceeds the max bson object size of ${maxBsonObjectSize}`
+    );
   }
 }
 
@@ -105,13 +248,15 @@ interface ClientInsertOperation {
  * @returns the operation.
  */
 export const buildInsertOneOperation = (
-  model: ClientInsertOneModel,
-  index: number
+  model: ClientInsertOneModel<Document>,
+  index: number,
+  pkFactory: PkFactory
 ): ClientInsertOperation => {
   const document: ClientInsertOperation = {
     insert: index,
     document: model.document
   };
+  document.document._id = model.document._id ?? pkFactory.createPk();
   return document;
 };
 
@@ -130,7 +275,10 @@ export interface ClientDeleteOperation {
  * @param index - The namespace index.
  * @returns the operation.
  */
-export const buildDeleteOneOperation = (model: ClientDeleteOneModel, index: number): Document => {
+export const buildDeleteOneOperation = (
+  model: ClientDeleteOneModel<Document>,
+  index: number
+): Document => {
   return createDeleteOperation(model, index, false);
 };
 
@@ -140,7 +288,10 @@ export const buildDeleteOneOperation = (model: ClientDeleteOneModel, index: numb
  * @param index - The namespace index.
  * @returns the operation.
  */
-export const buildDeleteManyOperation = (model: ClientDeleteManyModel, index: number): Document => {
+export const buildDeleteManyOperation = (
+  model: ClientDeleteManyModel<Document>,
+  index: number
+): Document => {
   return createDeleteOperation(model, index, true);
 };
 
@@ -148,7 +299,7 @@ export const buildDeleteManyOperation = (model: ClientDeleteManyModel, index: nu
  * Creates a delete operation based on the parameters.
  */
 function createDeleteOperation(
-  model: ClientDeleteOneModel | ClientDeleteManyModel,
+  model: ClientDeleteOneModel<Document> | ClientDeleteManyModel<Document>,
   index: number,
   multi: boolean
 ): ClientDeleteOperation {
@@ -175,6 +326,7 @@ export interface ClientUpdateOperation {
   hint?: Hint;
   upsert?: boolean;
   arrayFilters?: Document[];
+  collation?: CollationOptions;
 }
 
 /**
@@ -184,7 +336,7 @@ export interface ClientUpdateOperation {
  * @returns the operation.
  */
 export const buildUpdateOneOperation = (
-  model: ClientUpdateOneModel,
+  model: ClientUpdateOneModel<Document>,
   index: number
 ): ClientUpdateOperation => {
   return createUpdateOperation(model, index, false);
@@ -197,20 +349,37 @@ export const buildUpdateOneOperation = (
  * @returns the operation.
  */
 export const buildUpdateManyOperation = (
-  model: ClientUpdateManyModel,
+  model: ClientUpdateManyModel<Document>,
   index: number
 ): ClientUpdateOperation => {
   return createUpdateOperation(model, index, true);
 };
 
+/**
+ * Validate the update document.
+ * @param update - The update document.
+ */
+function validateUpdate(update: Document) {
+  if (!hasAtomicOperators(update)) {
+    throw new MongoAPIError(
+      'Client bulk write update models must only contain atomic modifiers (start with $) and must not be empty.'
+    );
+  }
+}
+
 /**
  * Creates a delete operation based on the parameters.
  */
 function createUpdateOperation(
-  model: ClientUpdateOneModel | ClientUpdateManyModel,
+  model: ClientUpdateOneModel<Document> | ClientUpdateManyModel<Document>,
   index: number,
   multi: boolean
 ): ClientUpdateOperation {
+  // Update documents provided in UpdateOne and UpdateMany write models are
+  // required only to contain atomic modifiers (i.e. keys that start with "$").
+  // Drivers MUST throw an error if an update document is empty or if the
+  // document's first key does not start with "$".
+  validateUpdate(model.update);
   const document: ClientUpdateOperation = {
     update: index,
     multi: multi,
@@ -226,6 +395,9 @@ function createUpdateOperation(
   if (model.arrayFilters) {
     document.arrayFilters = model.arrayFilters;
   }
+  if (model.collation) {
+    document.collation = model.collation;
+  }
   return document;
 }
 
@@ -237,6 +409,7 @@ export interface ClientReplaceOneOperation {
   updateMods: WithoutId<Document>;
   hint?: Hint;
   upsert?: boolean;
+  collation?: CollationOptions;
 }
 
 /**
@@ -246,9 +419,15 @@ export interface ClientReplaceOneOperation {
  * @returns the operation.
  */
 export const buildReplaceOneOperation = (
-  model: ClientReplaceOneModel,
+  model: ClientReplaceOneModel<Document>,
   index: number
 ): ClientReplaceOneOperation => {
+  if (hasAtomicOperators(model.replacement)) {
+    throw new MongoAPIError(
+      'Client bulk write replace models must not contain atomic modifiers (start with $) and must not be empty.'
+    );
+  }
+
   const document: ClientReplaceOneOperation = {
     update: index,
     multi: false,
@@ -261,14 +440,21 @@ export const buildReplaceOneOperation = (
   if (model.upsert) {
     document.upsert = model.upsert;
   }
+  if (model.collation) {
+    document.collation = model.collation;
+  }
   return document;
 };
 
 /** @internal */
-export function buildOperation(model: AnyClientBulkWriteModel, index: number): Document {
+export function buildOperation(
+  model: AnyClientBulkWriteModel<Document>,
+  index: number,
+  pkFactory: PkFactory
+): Document {
   switch (model.name) {
     case 'insertOne':
-      return buildInsertOneOperation(model, index);
+      return buildInsertOneOperation(model, index, pkFactory);
     case 'deleteOne':
       return buildDeleteOneOperation(model, index);
     case 'deleteMany':