WIP

Author: W-A-James

src/cmap/connection_pool.ts

@@ -27,7 +27,14 @@ import {
 } from '../error';
 import { CancellationToken, TypedEventEmitter } from '../mongo_types';
 import type { Server } from '../sdam/server';
-import { type Callback, eachAsync, List, makeCounter, TimeoutController } from '../utils';
+import {
+  type Callback,
+  eachAsync,
+  List,
+  makeCounter,
+  promiseWithResolvers,
+  TimeoutController
+} from '../utils';
 import { connect } from './connect';
 import { Connection, type ConnectionEvents, type ConnectionOptions } from './connection';
 import {
@@ -645,6 +652,56 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
     );
   }
 
+  async withConnectionAsync(
+    conn: Connection | undefined,
+    fn: AsyncWithConnectionCallback
+  ): Promise<Connection> {
+    let result;
+    if (conn) {
+      // use the provided connection, and do _not_ check it in after execution
+      try {
+        result = await fn(undefined, conn);
+        return result;
+      } catch (fnErr) {
+        if (fnErr instanceof MongoError && fnErr.code === MONGODB_ERROR_CODES.Reauthenticate) {
+          conn = this.reauthenticateAsync(conn);
+        } else {
+          throw fnErr;
+        }
+      }
+    } else {
+      const { promise, resolve, reject } = promiseWithResolvers<Connection>();
+      this.checkOut((err, conn) => {
+        fn(err as MongoError, conn).then(
+          () => {
+            if (conn) this.checkIn(conn);
+          },
+          fnErr => {
+            if (conn) {
+              this.withReauthenticationAsync(fnErr, conn, fn).then(resolve, reject);
+            } else {
+              reject(fnErr);
+            }
+          }
+        );
+      });
+      return promise;
+    }
+  }
+
+  async withReauthenticationAsync(
+    fnErr: AnyError,
+    conn: Connection,
+    fn: AsyncWithConnectionCallback
+  ): Promise<Connection> {
+    if (fnErr instanceof MongoError && fnErr.code === MONGODB_ERROR_CODES.Reauthenticate) {
+      conn = await this.reauthenticateAsync(conn);
+      return fn(undefined, conn);
+    } else {
+      throw fnErr;
+    }
+  }
+
   async reauthenticateAsync(connection: Connection): Promise<Connection> {
     const authContext = connection.authContext;
     if (!authContext) {
@@ -667,6 +724,7 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
     }
 
     await provider.reauth(authContext);
+
     return connection;
   }
 
@@ -945,3 +1003,8 @@ export type WithConnectionCallback = (
   connection: Connection | undefined,
   callback: Callback<Connection>
 ) => void;
+
+type AsyncWithConnectionCallback = (
+  error: MongoError | undefined,
+  connection: Connection | undefined
+) => Promise<Connection>;

src/operations/command.ts

@@ -152,6 +152,6 @@ export abstract class CommandOperation<T> extends AbstractOperation<T> {
       cmd = decorateWithExplain(cmd, this.explain);
     }
 
-    return server.commandAsync(this.ns, cmd, options);
+    return server.command(this.ns, cmd, options);
   }
 }

src/operations/find.ts

@@ -111,7 +111,7 @@ export class FindOperation extends CommandOperation<Document> {
       findCommand = decorateWithExplain(findCommand, this.explain);
     }
 
-    return server.commandAsync(this.ns, findCommand, {
+    return server.command(this.ns, findCommand, {
       ...this.options,
       ...this.bsonOptions,
       documentsReturnedIn: 'firstBatch',

src/operations/get_more.ts

@@ -96,7 +96,7 @@ export class GetMoreOperation extends AbstractOperation {
       ...this.options
     };
 
-    return server.commandAsync(this.ns, getMoreCmd, commandOptions);
+    return server.command(this.ns, getMoreCmd, commandOptions);
   }
 }
 

src/operations/kill_cursors.ts

@@ -46,7 +46,7 @@ export class KillCursorsOperation extends AbstractOperation {
       cursors: [this.cursorId]
     };
     try {
-      await server.commandAsync(this.ns, killCursorsCommand, { session });
+      await server.command(this.ns, killCursorsCommand, { session });
     } catch {
       // The driver should never emit errors from killCursors, this is spec-ed behavior
     }

src/operations/run_command.ts

@@ -28,7 +28,7 @@ export class RunCommandOperation<T = Document> extends AbstractOperation<T> {
 
   override async execute(server: Server, session: ClientSession | undefined): Promise<T> {
     this.server = server;
-    return server.commandAsync(this.ns, this.command, {
+    return server.command(this.ns, this.command, {
       ...this.options,
       readPreference: this.readPreference,
       session
@@ -54,7 +54,7 @@ export class RunAdminCommandOperation<T = Document> extends AbstractOperation<T>
 
   override async execute(server: Server, session: ClientSession | undefined): Promise<T> {
     this.server = server;
-    return server.commandAsync(this.ns, this.command, {
+    return server.command(this.ns, this.command, {
       ...this.options,
       readPreference: this.readPreference,
       session

src/operations/search_indexes/create.ts

@@ -36,7 +36,7 @@ export class CreateSearchIndexesOperation extends AbstractOperation<string[]> {
       indexes: this.descriptions
     };
 
-    const res = await server.commandAsync(namespace, command, { session });
+    const res = await server.command(namespace, command, { session });
 
     const indexesCreated: Array<{ name: string }> = res?.indexesCreated ?? [];
     return indexesCreated.map(({ name }) => name);

src/operations/search_indexes/drop.ts

@@ -28,7 +28,7 @@ export class DropSearchIndexOperation extends AbstractOperation<void> {
     }
 
     try {
-      await server.commandAsync(namespace, command, { session });
+      await server.command(namespace, command, { session });
     } catch (error) {
       const isNamespaceNotFoundError =
         error instanceof MongoServerError && error.code === MONGODB_ERROR_CODES.NamespaceNotFound;

src/operations/search_indexes/update.ts

@@ -27,7 +27,7 @@ export class UpdateSearchIndexOperation extends AbstractOperation<void> {
       definition: this.definition
     };
 
-    await server.commandAsync(namespace, command, { session });
+    await server.command(namespace, command, { session });
     return;
   }
 }