From 546d2b28afd2112369499f211e4d3445c70636e2 Mon Sep 17 00:00:00 2001
From: Warren James <warren.james@mongodb.com>
Date: Wed, 25 Sep 2024 17:43:12 -0400
Subject: [PATCH] chore: allow clientBulkWrite to use TimeoutContext (#4251)

---
 .../client_bulk_write/client_bulk_write.ts         | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/src/operations/client_bulk_write/client_bulk_write.ts b/src/operations/client_bulk_write/client_bulk_write.ts
index e901407cd78..26d1e7bb60f 100644
--- a/src/operations/client_bulk_write/client_bulk_write.ts
+++ b/src/operations/client_bulk_write/client_bulk_write.ts
@@ -2,6 +2,7 @@ import { MongoClientBulkWriteExecutionError, ServerType } from '../../beta';
 import { ClientBulkWriteCursorResponse } from '../../cmap/wire_protocol/responses';
 import type { Server } from '../../sdam/server';
 import type { ClientSession } from '../../sessions';
+import { type TimeoutContext } from '../../timeout';
 import { MongoDBNamespace } from '../../utils';
 import { CommandOperation } from '../command';
 import { Aspect, defineAspects } from '../operation';
@@ -43,7 +44,8 @@ export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCu
    */
   override async execute(
     server: Server,
-    session: ClientSession | undefined
+    session: ClientSession | undefined,
+    timeoutContext: TimeoutContext
   ): Promise<ClientBulkWriteCursorResponse> {
     let command;
 
@@ -52,7 +54,7 @@ export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCu
         let connection;
         if (!session.pinnedConnection) {
           // Checkout a connection to build the command.
-          connection = await server.pool.checkOut();
+          connection = await server.pool.checkOut({ timeoutContext });
           // 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);
@@ -93,7 +95,13 @@ export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCu
     if (!this.canRetryWrite) {
       this.options.willRetryWrite = false;
     }
-    return await super.executeCommand(server, session, command, ClientBulkWriteCursorResponse);
+    return await super.executeCommand(
+      server,
+      session,
+      command,
+      timeoutContext,
+      ClientBulkWriteCursorResponse
+    );
   }
 }