mongodb 6.9.0-dev.20241001.sha.85f7dcf9 → 6.9.0-dev.20241003.sha.91f30357

package/src/cmap/commands.ts

@@ -429,10 +429,60 @@ export interface OpMsgOptions {
 
 /** @internal */
 export class DocumentSequence {
+  field: string;
   documents: Document[];
+  serializedDocumentsLength: number;
+  private chunks: Uint8Array[];
+  private header: Buffer;
 
-  constructor(documents: Document[]) {
-    this.documents = documents;
+  /**
+   * Create a new document sequence for the provided field.
+   * @param field - The field it will replace.
+   */
+  constructor(field: string, documents?: Document[]) {
+    this.field = field;
+    this.documents = [];
+    this.chunks = [];
+    this.serializedDocumentsLength = 0;
+    // Document sequences starts with type 1 at the first byte.
+    // Field strings must always be UTF-8.
+    const buffer = Buffer.allocUnsafe(1 + 4 + this.field.length + 1);
+    buffer[0] = 1;
+    // Third part is the field name at offset 5 with trailing null byte.
+    encodeUTF8Into(buffer, `${this.field}\0`, 5);
+    this.chunks.push(buffer);
+    this.header = buffer;
+    if (documents) {
+      for (const doc of documents) {
+        this.push(doc, BSON.serialize(doc));
+      }
+    }
+  }
+
+  /**
+   * Push a document to the document sequence. Will serialize the document
+   * as well and return the current serialized length of all documents.
+   * @param document - The document to add.
+   * @param buffer - The serialized document in raw BSON.
+   * @returns The new total document sequence length.
+   */
+  push(document: Document, buffer: Uint8Array): number {
+    this.serializedDocumentsLength += buffer.length;
+    // Push the document.
+    this.documents.push(document);
+    // Push the document raw bson.
+    this.chunks.push(buffer);
+    // Write the new length.
+    this.header?.writeInt32LE(4 + this.field.length + 1 + this.serializedDocumentsLength, 1);
+    return this.serializedDocumentsLength + this.header.length;
+  }
+
+  /**
+   * Get the fully serialized bytes for the document sequence section.
+   * @returns The section bytes.
+   */
+  toBin(): Uint8Array {
+    return Buffer.concat(this.chunks);
   }
 }
 
@@ -543,21 +593,7 @@ export class OpMsgRequest {
     const chunks = [];
     for (const [key, value] of Object.entries(document)) {
       if (value instanceof DocumentSequence) {
-        // Document sequences starts with type 1 at the first byte.
-        const buffer = Buffer.allocUnsafe(1 + 4 + key.length + 1);
-        buffer[0] = 1;
-        // Third part is the field name at offset 5 with trailing null byte.
-        encodeUTF8Into(buffer, `${key}\0`, 5);
-        chunks.push(buffer);
-        // Fourth part are the documents' bytes.
-        let docsLength = 0;
-        for (const doc of value.documents) {
-          const docBson = this.serializeBson(doc);
-          docsLength += docBson.length;
-          chunks.push(docBson);
-        }
-        // Second part of the sequence is the length at offset 1;
-        buffer.writeInt32LE(4 + key.length + 1 + docsLength, 1);
+        chunks.push(value.toBin());
         // Why are we removing the field from the command? This is because it needs to be
         // removed in the OP_MSG request first section, and DocumentSequence is not a
         // BSON type and is specific to the MongoDB wire protocol so there's nothing

package/src/cursor/client_bulk_write_cursor.ts

@@ -1,8 +1,10 @@
-import type { Document } from '../bson';
+import { type Document } from 'bson';
+
 import { type ClientBulkWriteCursorResponse } from '../cmap/wire_protocol/responses';
-import { MongoBulkWriteCursorError } from '../error';
+import { MongoClientBulkWriteCursorError } from '../error';
 import type { MongoClient } from '../mongo_client';
 import { ClientBulkWriteOperation } from '../operations/client_bulk_write/client_bulk_write';
+import { type ClientBulkWriteCommandBuilder } from '../operations/client_bulk_write/command_builder';
 import { type ClientBulkWriteOptions } from '../operations/client_bulk_write/common';
 import { executeOperation } from '../operations/execute_operation';
 import type { ClientSession } from '../sessions';
@@ -24,17 +26,21 @@ export interface ClientBulkWriteCursorOptions
  * @internal
  */
 export class ClientBulkWriteCursor extends AbstractCursor {
-  public readonly command: Document;
+  commandBuilder: ClientBulkWriteCommandBuilder;
   /** @internal */
   private cursorResponse?: ClientBulkWriteCursorResponse;
   /** @internal */
   private clientBulkWriteOptions: ClientBulkWriteOptions;
 
   /** @internal */
-  constructor(client: MongoClient, command: Document, options: ClientBulkWriteOptions = {}) {
+  constructor(
+    client: MongoClient,
+    commandBuilder: ClientBulkWriteCommandBuilder,
+    options: ClientBulkWriteOptions = {}
+  ) {
     super(client, new MongoDBNamespace('admin', '$cmd'), options);
 
-    this.command = command;
+    this.commandBuilder = commandBuilder;
     this.clientBulkWriteOptions = options;
   }
 
@@ -44,22 +50,29 @@ export class ClientBulkWriteCursor extends AbstractCursor {
    */
   get response(): ClientBulkWriteCursorResponse {
     if (this.cursorResponse) return this.cursorResponse;
-    throw new MongoBulkWriteCursorError(
+    throw new MongoClientBulkWriteCursorError(
       'No client bulk write cursor response returned from the server.'
     );
   }
 
+  /**
+   * Get the last set of operations the cursor executed.
+   */
+  get operations(): Document[] {
+    return this.commandBuilder.lastOperations;
+  }
+
   clone(): ClientBulkWriteCursor {
     const clonedOptions = mergeOptions({}, this.clientBulkWriteOptions);
     delete clonedOptions.session;
-    return new ClientBulkWriteCursor(this.client, this.command, {
+    return new ClientBulkWriteCursor(this.client, this.commandBuilder, {
       ...clonedOptions
     });
   }
 
   /** @internal */
   async _initialize(session: ClientSession): Promise<InitialCursorResponse> {
-    const clientBulkWriteOperation = new ClientBulkWriteOperation(this.command, {
+    const clientBulkWriteOperation = new ClientBulkWriteOperation(this.commandBuilder, {
       ...this.clientBulkWriteOptions,
       ...this.cursorOptions,
       session

package/src/error.ts

@@ -622,7 +622,7 @@ export class MongoGCPError extends MongoOIDCError {
  * @public
  * @category Error
  */
-export class MongoBulkWriteCursorError extends MongoRuntimeError {
+export class MongoClientBulkWriteCursorError extends MongoRuntimeError {
   /**
    * **Do not use this constructor!**
    *
@@ -639,7 +639,34 @@ export class MongoBulkWriteCursorError extends MongoRuntimeError {
   }
 
   override get name(): string {
-    return 'MongoBulkWriteCursorError';
+    return 'MongoClientBulkWriteCursorError';
+  }
+}
+
+/**
+ * An error indicating that an error occurred on the client when executing a client bulk write.
+ *
+ * @public
+ * @category Error
+ */
+export class MongoClientBulkWriteExecutionError extends MongoRuntimeError {
+  /**
+   * **Do not use this constructor!**
+   *
+   * Meant for internal use only.
+   *
+   * @remarks
+   * This class is only meant to be constructed within the driver. This constructor is
+   * not subject to semantic versioning compatibility guarantees and may change at any time.
+   *
+   * @public
+   **/
+  constructor(message: string) {
+    super(message);
+  }
+
+  override get name(): string {
+    return 'MongoClientBulkWriteExecutionError';
   }
 }
 

package/src/index.ts

@@ -44,8 +44,9 @@ export {
   MongoAWSError,
   MongoAzureError,
   MongoBatchReExecutionError,
-  MongoBulkWriteCursorError,
   MongoChangeStreamError,
+  MongoClientBulkWriteCursorError,
+  MongoClientBulkWriteExecutionError,
   MongoCompatibilityError,
   MongoCursorExhaustedError,
   MongoCursorInUseError,

package/src/operations/client_bulk_write/client_bulk_write.ts

@@ -1,11 +1,11 @@
-import { type Document } from 'bson';
-
+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';
 
 /**
@@ -13,16 +13,16 @@ import { type ClientBulkWriteOptions } from './common';
  * @internal
  */
 export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCursorResponse> {
-  command: Document;
+  commandBuilder: ClientBulkWriteCommandBuilder;
   override options: ClientBulkWriteOptions;
 
   override get commandName() {
     return 'bulkWrite' as const;
   }
 
-  constructor(command: Document, options: ClientBulkWriteOptions) {
+  constructor(commandBuilder: ClientBulkWriteCommandBuilder, options: ClientBulkWriteOptions) {
     super(undefined, options);
-    this.command = command;
+    this.commandBuilder = commandBuilder;
     this.options = options;
     this.ns = new MongoDBNamespace('admin', '$cmd');
   }
@@ -37,9 +37,45 @@ export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCu
     server: Server,
     session: ClientSession | undefined
   ): Promise<ClientBulkWriteCursorResponse> {
-    return await super.executeCommand(server, session, this.command, ClientBulkWriteCursorResponse);
+    let command;
+
+    if (server.description.type === ServerType.LoadBalancer) {
+      if (session) {
+        // Checkout a connection to build the command.
+        const 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);
+        command = this.commandBuilder.buildBatch(
+          connection.hello?.maxMessageSizeBytes,
+          connection.hello?.maxWriteBatchSize
+        );
+      } 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) {
+        throw new MongoClientBulkWriteExecutionError(
+          'In order to execute a client bulk write, both maxWriteBatchSize and maxMessageSizeBytes must be provided by the servers hello response.'
+        );
+      }
+      command = this.commandBuilder.buildBatch(
+        server.description.maxMessageSizeBytes,
+        server.description.maxWriteBatchSize
+      );
+    }
+    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]);
+defineAspects(ClientBulkWriteOperation, [
+  Aspect.WRITE_OPERATION,
+  Aspect.SKIP_COLLATION,
+  Aspect.CURSOR_CREATING
+]);

package/src/operations/client_bulk_write/command_builder.ts

@@ -1,4 +1,4 @@
-import { type Document } from '../../bson';
+import { BSON, type Document } from '../../bson';
 import { DocumentSequence } from '../../cmap/commands';
 import { type PkFactory } from '../../mongo_client';
 import type { Filter, OptionalId, UpdateFilter, WithoutId } from '../../mongo_types';
@@ -28,11 +28,18 @@ export interface ClientBulkWriteCommand {
   comment?: any;
 }
 
+/**
+ * The bytes overhead for the extra fields added post command generation.
+ */
+const MESSAGE_OVERHEAD_BYTES = 1000;
+
 /** @internal */
 export class ClientBulkWriteCommandBuilder {
   models: AnyClientBulkWriteModel[];
   options: ClientBulkWriteOptions;
   pkFactory: PkFactory;
+  currentModelIndex: number;
+  lastOperations: Document[];
 
   /**
    * Create the command builder.
@@ -46,6 +53,8 @@ export class ClientBulkWriteCommandBuilder {
     this.models = models;
     this.options = options;
     this.pkFactory = pkFactory ?? DEFAULT_PK_FACTORY;
+    this.currentModelIndex = 0;
+    this.lastOperations = [];
   }
 
   /**
@@ -60,34 +69,96 @@ 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;
+  }
+
+  /**
+   * 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): ClientBulkWriteCommand {
+    let commandLength = 0;
     let currentNamespaceIndex = 0;
+    const command: ClientBulkWriteCommand = this.baseCommand();
     const namespaces = new Map<string, number>();
-    for (const model of this.models) {
+
+    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, this.pkFactory));
+      const nsIndex = namespaces.get(ns);
+
+      if (nsIndex != null) {
+        // Build the operation and serialize it to get the bytes buffer.
+        const operation = buildOperation(model, nsIndex, this.pkFactory);
+        const operationBuffer = BSON.serialize(operation);
+
+        // 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, this.pkFactory));
-        currentNamespaceIndex++;
+        const nsInfo = { ns: ns };
+        const nsInfoBuffer = BSON.serialize(nsInfo);
+        const operation = buildOperation(model, currentNamespaceIndex, this.pkFactory);
+        const operationBuffer = BSON.serialize(operation);
+
+        // 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) {
@@ -103,7 +174,8 @@ export class ClientBulkWriteCommandBuilder {
     if (this.options.comment !== undefined) {
       command.comment = this.options.comment;
     }
-    return [command];
+
+    return command;
   }
 }
 

package/src/operations/client_bulk_write/executor.ts

@@ -1,11 +1,9 @@
-import { type Document } from 'bson';
-
 import { ClientBulkWriteCursor } from '../../cursor/client_bulk_write_cursor';
 import { type MongoClient } from '../../mongo_client';
 import { WriteConcern } from '../../write_concern';
 import { executeOperation } from '../execute_operation';
 import { ClientBulkWriteOperation } from './client_bulk_write';
-import { type ClientBulkWriteCommand, ClientBulkWriteCommandBuilder } from './command_builder';
+import { ClientBulkWriteCommandBuilder } from './command_builder';
 import {
   type AnyClientBulkWriteModel,
   type ClientBulkWriteOptions,
@@ -57,43 +55,26 @@ export class ClientBulkWriteExecutor {
       this.options,
       pkFactory
     );
-    const commands = commandBuilder.buildCommands();
+    // Unacknowledged writes need to execute all batches and return { ok: 1}
     if (this.options.writeConcern?.w === 0) {
-      return await executeUnacknowledged(this.client, this.options, commands);
+      while (commandBuilder.hasNextBatch()) {
+        const operation = new ClientBulkWriteOperation(commandBuilder, this.options);
+        await executeOperation(this.client, operation);
+      }
+      return { ok: 1 };
+    } else {
+      const resultsMerger = new ClientBulkWriteResultsMerger(this.options);
+      // For each command will will create and exhaust a cursor for the results.
+      let currentBatchOffset = 0;
+      while (commandBuilder.hasNextBatch()) {
+        const cursor = new ClientBulkWriteCursor(this.client, commandBuilder, this.options);
+        const docs = await cursor.toArray();
+        const operations = cursor.operations;
+        resultsMerger.merge(currentBatchOffset, operations, cursor.response, docs);
+        // Set the new batch index so we can back back to the index in the original models.
+        currentBatchOffset += operations.length;
+      }
+      return resultsMerger.result;
     }
-    return await executeAcknowledged(this.client, this.options, commands);
-  }
-}
-
-/**
- * Execute an acknowledged bulk write.
- */
-async function executeAcknowledged(
-  client: MongoClient,
-  options: ClientBulkWriteOptions,
-  commands: ClientBulkWriteCommand[]
-): Promise<ClientBulkWriteResult> {
-  const resultsMerger = new ClientBulkWriteResultsMerger(options);
-  // For each command will will create and exhaust a cursor for the results.
-  for (const command of commands) {
-    const cursor = new ClientBulkWriteCursor(client, command, options);
-    const docs = await cursor.toArray();
-    resultsMerger.merge(command.ops.documents, cursor.response, docs);
-  }
-  return resultsMerger.result;
-}
-
-/**
- * Execute an unacknowledged bulk write.
- */
-async function executeUnacknowledged(
-  client: MongoClient,
-  options: ClientBulkWriteOptions,
-  commands: Document[]
-): Promise<{ ok: 1 }> {
-  for (const command of commands) {
-    const operation = new ClientBulkWriteOperation(command, options);
-    await executeOperation(client, operation);
   }
-  return { ok: 1 };
 }

package/src/operations/client_bulk_write/results_merger.ts

@@ -42,11 +42,13 @@ export class ClientBulkWriteResultsMerger {
 
   /**
    * Merge the results in the cursor to the existing result.
+   * @param currentBatchOffset - The offset index to the original models.
    * @param response - The cursor response.
    * @param documents - The documents in the cursor.
    * @returns The current result.
    */
   merge(
+    currentBatchOffset: number,
     operations: Document[],
     response: ClientBulkWriteCursorResponse,
     documents: Document[]
@@ -67,7 +69,9 @@ export class ClientBulkWriteResultsMerger {
           const operation = operations[document.idx];
           // Handle insert results.
           if ('insert' in operation) {
-            this.result.insertResults?.set(document.idx, { insertedId: operation.document._id });
+            this.result.insertResults?.set(document.idx + currentBatchOffset, {
+              insertedId: operation.document._id
+            });
           }
           // Handle update results.
           if ('update' in operation) {
@@ -80,11 +84,13 @@ export class ClientBulkWriteResultsMerger {
             if (document.upserted) {
               result.upsertedId = document.upserted._id;
             }
-            this.result.updateResults?.set(document.idx, result);
+            this.result.updateResults?.set(document.idx + currentBatchOffset, result);
           }
           // Handle delete results.
           if ('delete' in operation) {
-            this.result.deleteResults?.set(document.idx, { deletedCount: document.n });
+            this.result.deleteResults?.set(document.idx + currentBatchOffset, {
+              deletedCount: document.n
+            });
           }
         }
       }

package/src/sdam/server.ts

@@ -513,7 +513,8 @@ function isPinnableCommand(cmd: Document, session?: ClientSession): boolean {
       'find' in cmd ||
       'getMore' in cmd ||
       'listCollections' in cmd ||
-      'listIndexes' in cmd
+      'listIndexes' in cmd ||
+      'bulkWrite' in cmd
     );
   }
 

package/src/sdam/server_description.ts

@@ -69,6 +69,12 @@ export class ServerDescription {
   setVersion: number | null;
   electionId: ObjectId | null;
   logicalSessionTimeoutMinutes: number | null;
+  /** The max message size in bytes for the server. */
+  maxMessageSizeBytes: number | null;
+  /** The max number of writes in a bulk write command. */
+  maxWriteBatchSize: number | null;
+  /** The max bson object size. */
+  maxBsonObjectSize: number | null;
 
   // NOTE: does this belong here? It seems we should gossip the cluster time at the CMAP level
   $clusterTime?: ClusterTime;
@@ -111,6 +117,9 @@ export class ServerDescription {
     this.setVersion = hello?.setVersion ?? null;
     this.electionId = hello?.electionId ?? null;
     this.logicalSessionTimeoutMinutes = hello?.logicalSessionTimeoutMinutes ?? null;
+    this.maxMessageSizeBytes = hello?.maxMessageSizeBytes ?? null;
+    this.maxWriteBatchSize = hello?.maxWriteBatchSize ?? null;
+    this.maxBsonObjectSize = hello?.maxBsonObjectSize ?? null;
     this.primary = hello?.primary ?? null;
     this.me = hello?.me?.toLowerCase() ?? null;
     this.$clusterTime = hello?.$clusterTime ?? null;

package/src/sdam/topology_description.ts

@@ -43,7 +43,6 @@ export class TopologyDescription {
   heartbeatFrequencyMS: number;
   localThresholdMS: number;
   commonWireVersion: number;
-
   /**
    * Create a TopologyDescription
    */