feat(NODE-6392): add timeoutMS support to ClientEncryption helpers part 1 (mongodb#4281)

Author: nbbeeken

src/client-side-encryption/auto_encrypter.ts

@@ -310,7 +310,10 @@ export class AutoEncrypter {
       // eslint-disable-next-line @typescript-eslint/ban-ts-comment
       // @ts-ignore: TS complains as this always returns true on versions where it is present.
       if (net.getDefaultAutoSelectFamily) {
-        Object.assign(clientOptions, autoSelectSocketOptions(this._client.options));
+        // AutoEncrypter is made inside of MongoClient constructor while options are being parsed,
+        // we do not have access to the options that are in progress.
+        // TODO(NODE-NODE-6449): AutoEncrypter does not use client options for autoSelectFamily
+        Object.assign(clientOptions, autoSelectSocketOptions(this._client.s?.options ?? {}));
       }
 
       this._mongocryptdClient = new MongoClient(this._mongocryptdManager.uri, clientOptions);
@@ -392,7 +395,7 @@ export class AutoEncrypter {
       promoteLongs: false,
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
     return deserialize(await stateMachine.execute(this, context, options.timeoutContext), {
@@ -413,7 +416,7 @@ export class AutoEncrypter {
       ...options,
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
     return await stateMachine.execute(

src/client-side-encryption/client_encryption.ts

@@ -24,7 +24,8 @@ import { type MongoClient, type MongoClientOptions } from '../mongo_client';
 import { type Filter, type WithId } from '../mongo_types';
 import { type CreateCollectionOptions } from '../operations/create_collection';
 import { type DeleteResult } from '../operations/delete';
-import { MongoDBCollectionNamespace } from '../utils';
+import { TimeoutContext } from '../timeout';
+import { MongoDBCollectionNamespace, resolveTimeoutOptions } from '../utils';
 import * as cryptoCallbacks from './crypto_callbacks';
 import {
   MongoCryptCreateDataKeyError,
@@ -74,6 +75,8 @@ export class ClientEncryption {
   _tlsOptions: CSFLEKMSTlsOptions;
   /** @internal */
   _kmsProviders: KMSProviders;
+  /** @internal */
+  _timeoutMS?: number;
 
   /** @internal */
   _mongoCrypt: MongoCrypt;
@@ -120,6 +123,8 @@ export class ClientEncryption {
     this._proxyOptions = options.proxyOptions ?? {};
     this._tlsOptions = options.tlsOptions ?? {};
     this._kmsProviders = options.kmsProviders || {};
+    const { timeoutMS } = resolveTimeoutOptions(client, options);
+    this._timeoutMS = timeoutMS;
 
     if (options.keyVaultNamespace == null) {
       throw new MongoCryptInvalidArgumentError('Missing required option `keyVaultNamespace`');
@@ -212,7 +217,7 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
     const dataKey = deserialize(await stateMachine.execute(this, context)) as DataKey;
@@ -270,10 +275,14 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
-    const { v: dataKeys } = deserialize(await stateMachine.execute(this, context));
+    const timeoutContext = TimeoutContext.create(
+      resolveTimeoutOptions(this._client, { timeoutMS: this._timeoutMS })
+    );
+
+    const { v: dataKeys } = deserialize(await stateMachine.execute(this, context, timeoutContext));
     if (dataKeys.length === 0) {
       return {};
     }
@@ -303,7 +312,8 @@ export class ClientEncryption {
       .db(dbName)
       .collection<DataKey>(collectionName)
       .bulkWrite(replacements, {
-        writeConcern: { w: 'majority' }
+        writeConcern: { w: 'majority' },
+        timeoutMS: timeoutContext.csotEnabled() ? timeoutContext?.remainingTimeMS : undefined
       });
 
     return { bulkWriteResult: result };
@@ -332,7 +342,7 @@ export class ClientEncryption {
     return await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
-      .deleteOne({ _id }, { writeConcern: { w: 'majority' } });
+      .deleteOne({ _id }, { writeConcern: { w: 'majority' }, timeoutMS: this._timeoutMS });
   }
 
   /**
@@ -355,7 +365,7 @@ export class ClientEncryption {
     return this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
-      .find({}, { readConcern: { level: 'majority' } });
+      .find({}, { readConcern: { level: 'majority' }, timeoutMS: this._timeoutMS });
   }
 
   /**
@@ -381,7 +391,7 @@ export class ClientEncryption {
     return await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
-      .findOne({ _id }, { readConcern: { level: 'majority' } });
+      .findOne({ _id }, { readConcern: { level: 'majority' }, timeoutMS: this._timeoutMS });
   }
 
   /**
@@ -408,7 +418,10 @@ export class ClientEncryption {
     return await this._keyVaultClient
       .db(dbName)
       .collection<DataKey>(collectionName)
-      .findOne({ keyAltNames: keyAltName }, { readConcern: { level: 'majority' } });
+      .findOne(
+        { keyAltNames: keyAltName },
+        { readConcern: { level: 'majority' }, timeoutMS: this._timeoutMS }
+      );
   }
 
   /**
@@ -442,7 +455,7 @@ export class ClientEncryption {
       .findOneAndUpdate(
         { _id },
         { $addToSet: { keyAltNames: keyAltName } },
-        { writeConcern: { w: 'majority' }, returnDocument: 'before' }
+        { writeConcern: { w: 'majority' }, returnDocument: 'before', timeoutMS: this._timeoutMS }
       );
 
     return value;
@@ -503,7 +516,8 @@ export class ClientEncryption {
       .collection<DataKey>(collectionName)
       .findOneAndUpdate({ _id }, pipeline, {
         writeConcern: { w: 'majority' },
-        returnDocument: 'before'
+        returnDocument: 'before',
+        timeoutMS: this._timeoutMS
       });
 
     return value;
@@ -650,7 +664,7 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
 
     const { v } = deserialize(await stateMachine.execute(this, context));
@@ -729,7 +743,7 @@ export class ClientEncryption {
     const stateMachine = new StateMachine({
       proxyOptions: this._proxyOptions,
       tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.options)
+      socketOptions: autoSelectSocketOptions(this._client.s.options)
     });
     const context = this._mongoCrypt.makeExplicitEncryptionContext(valueBuffer, contextOptions);
 
@@ -818,6 +832,11 @@ export interface ClientEncryptionOptions {
    * TLS options for kms providers to use.
    */
   tlsOptions?: CSFLEKMSTlsOptions;
+
+  /**
+   * The timeout setting to be used for all the operations on ClientEncryption.
+   */
+  timeoutMS?: number;
 }
 
 /**

src/cursor/abstract_cursor.ts

@@ -835,7 +835,7 @@ export abstract class AbstractCursor<
     if (this.cursorOptions.timeoutMS != null) {
       this.timeoutContext ??= new CursorTimeoutContext(
         TimeoutContext.create({
-          serverSelectionTimeoutMS: this.client.options.serverSelectionTimeoutMS,
+          serverSelectionTimeoutMS: this.client.s.options.serverSelectionTimeoutMS,
           timeoutMS: this.cursorOptions.timeoutMS
         }),
         this
@@ -925,7 +925,7 @@ export abstract class AbstractCursor<
         this.timeoutContext?.clear();
         return new CursorTimeoutContext(
           TimeoutContext.create({
-            serverSelectionTimeoutMS: this.client.options.serverSelectionTimeoutMS,
+            serverSelectionTimeoutMS: this.client.s.options.serverSelectionTimeoutMS,
             timeoutMS
           }),
           this

src/gridfs/index.ts

@@ -161,7 +161,7 @@ export class GridFSBucket extends TypedEventEmitter<GridFSBucketEvents> {
     if (timeoutMS) {
       timeoutContext = new CSOTTimeoutContext({
         timeoutMS,
-        serverSelectionTimeoutMS: this.s.db.client.options.serverSelectionTimeoutMS
+        serverSelectionTimeoutMS: this.s.db.client.s.options.serverSelectionTimeoutMS
       });
     }
 
@@ -241,7 +241,7 @@ export class GridFSBucket extends TypedEventEmitter<GridFSBucketEvents> {
     if (timeoutMS) {
       timeoutContext = new CSOTTimeoutContext({
         timeoutMS,
-        serverSelectionTimeoutMS: this.s.db.client.options.serverSelectionTimeoutMS
+        serverSelectionTimeoutMS: this.s.db.client.s.options.serverSelectionTimeoutMS
       });
     }
 

src/gridfs/upload.ts

@@ -10,7 +10,7 @@ import {
   MongoOperationTimeoutError
 } from '../error';
 import { CSOTTimeoutContext } from '../timeout';
-import { type Callback, squashError } from '../utils';
+import { type Callback, resolveTimeoutOptions, squashError } from '../utils';
 import type { WriteConcernOptions } from '../write_concern';
 import { WriteConcern } from './../write_concern';
 import type { GridFSFile } from './download';
@@ -143,7 +143,8 @@ export class GridFSBucketWriteStream extends Writable {
     if (options.timeoutMS != null)
       this.timeoutContext = new CSOTTimeoutContext({
         timeoutMS: options.timeoutMS,
-        serverSelectionTimeoutMS: this.bucket.s.db.client.options.serverSelectionTimeoutMS
+        serverSelectionTimeoutMS: resolveTimeoutOptions(this.bucket.s.db.client, {})
+          .serverSelectionTimeoutMS
       });
   }
 

src/mongo_client.ts

@@ -490,7 +490,7 @@ export class MongoClient extends TypedEventEmitter<MongoClientEvents> implements
 
   /** @internal */
   get timeoutMS(): number | undefined {
-    return this.options.timeoutMS;
+    return this.s.options.timeoutMS;
   }
 
   /**
@@ -706,7 +706,7 @@ export class MongoClient extends TypedEventEmitter<MongoClientEvents> implements
 
     // Default to db from connection string if not provided
     if (!dbName) {
-      dbName = this.options.dbName;
+      dbName = this.s.options.dbName;
     }
 
     // Copy the options and add out internal override of the not shared flag

src/operations/client_bulk_write/executor.ts

@@ -56,7 +56,7 @@ export class ClientBulkWriteExecutor {
 
     // If no write concern was provided, we inherit one from the client.
     if (!this.options.writeConcern) {
-      this.options.writeConcern = WriteConcern.fromOptions(this.client.options);
+      this.options.writeConcern = WriteConcern.fromOptions(this.client.s.options);
     }
 
     if (this.options.writeConcern?.w === 0) {

src/operations/create_collection.ts

@@ -137,7 +137,7 @@ export class CreateCollectionOperation extends CommandOperation<Collection> {
 
     const encryptedFields: Document | undefined =
       options.encryptedFields ??
-      db.client.options.autoEncryption?.encryptedFieldsMap?.[`${db.databaseName}.${name}`];
+      db.client.s.options.autoEncryption?.encryptedFieldsMap?.[`${db.databaseName}.${name}`];
 
     if (encryptedFields) {
       // Creating a QE collection required min server of 7.0.0

src/operations/drop.ts

@@ -39,7 +39,7 @@ export class DropCollectionOperation extends CommandOperation<boolean> {
     const options = this.options;
     const name = this.name;
 
-    const encryptedFieldsMap = db.client.options.autoEncryption?.encryptedFieldsMap;
+    const encryptedFieldsMap = db.client.s.options.autoEncryption?.encryptedFieldsMap;
     let encryptedFields: Document | undefined =
       options.encryptedFields ?? encryptedFieldsMap?.[`${db.databaseName}.${name}`];