rebase

Author: baileympearson

drivers-evergreen-tools

@@ -10,10 +10,7 @@ import { MongoInvalidArgumentError } from './error';
 import type { MongoClient, PkFactory } from './mongo_client';
 import type { Abortable, TODO_NODE_3286 } from './mongo_types';
 import type { AggregateOptions } from './operations/aggregate';
-import {
-  CreateCollectionOperation,
-  type CreateCollectionOptions
-} from './operations/create_collection';
+import { type CreateCollectionOptions, createCollections } from './operations/create_collection';
 import {
   type DropCollectionOptions,
   dropCollections,
@@ -241,10 +238,15 @@ export class Db {
     name: string,
     options?: CreateCollectionOptions
   ): Promise<Collection<TSchema>> {
-    return await executeOperation(
-      this.client,
-      new CreateCollectionOperation(this, name, resolveOptions(this, options)) as TODO_NODE_3286
-    );
+    options = resolveOptions(this, { ...options });
+    if (options.session) {
+      return await createCollections<TSchema>(this, name, options);
+    }
+
+    return await this.client.withSession({ explicit: false }, async session => {
+      options.session = session;
+      return await createCollections<TSchema>(this, name, options);
+    });
   }
 
   /**

src/db.ts

@@ -9,8 +9,9 @@ import { MongoCompatibilityError } from '../error';
 import type { PkFactory } from '../mongo_client';
 import type { Server } from '../sdam/server';
 import type { ClientSession } from '../sessions';
-import { type TimeoutContext } from '../timeout';
+import { TimeoutContext } from '../timeout';
 import { CommandOperation, type CommandOperationOptions } from './command';
+import { executeOperation } from './execute_operation';
 import { CreateIndexesOperation } from './indexes';
 import { Aspect, defineAspects } from './operation';
 
@@ -102,6 +103,9 @@ export interface CreateCollectionOptions extends CommandOperationOptions {
    * change streams that listen on this collection.
    */
   changeStreamPreAndPostImages?: { enabled: boolean };
+
+  /** @internal */
+  isEncryptedCollection?: boolean;
 }
 
 /* @internal */
@@ -131,15 +135,7 @@ export class CreateCollectionOperation extends CommandOperation<Collection> {
     session: ClientSession | undefined,
     timeoutContext: TimeoutContext
   ): Promise<Collection> {
-    const db = this.db;
-    const name = this.name;
-    const options = this.options;
-
-    const encryptedFields: Document | undefined =
-      options.encryptedFields ??
-      db.client.s.options.autoEncryption?.encryptedFieldsMap?.[`${db.databaseName}.${name}`];
-
-    if (encryptedFields) {
+    if (this.options.isEncryptedCollection) {
       // Creating a QE collection required min server of 7.0.0
       // TODO(NODE-5353): Get wire version information from connection.
       if (
@@ -150,64 +146,81 @@ export class CreateCollectionOperation extends CommandOperation<Collection> {
           `${INVALID_QE_VERSION} The minimum server version required is ${MIN_SUPPORTED_QE_SERVER_VERSION}`
         );
       }
-      // Create auxilliary collections for queryable encryption support.
-      const escCollection = encryptedFields.escCollection ?? `enxcol_.${name}.esc`;
-      const ecocCollection = encryptedFields.ecocCollection ?? `enxcol_.${name}.ecoc`;
-
-      for (const collectionName of [escCollection, ecocCollection]) {
-        const createOp = new CreateCollectionOperation(db, collectionName, {
-          clusteredIndex: {
-            key: { _id: 1 },
-            unique: true
-          }
-        });
-        await createOp.executeWithoutEncryptedFieldsCheck(server, session, timeoutContext);
-      }
-
-      if (!options.encryptedFields) {
-        this.options = { ...this.options, encryptedFields };
-      }
-    }
-
-    const coll = await this.executeWithoutEncryptedFieldsCheck(server, session, timeoutContext);
-
-    if (encryptedFields) {
-      // Create the required index for queryable encryption support.
-      const createIndexOp = CreateIndexesOperation.fromIndexSpecification(
-        db,
-        name,
-        { __safeContent__: 1 },
-        {}
-      );
-      await createIndexOp.execute(server, session, timeoutContext);
     }
 
-    return coll;
-  }
-
-  private async executeWithoutEncryptedFieldsCheck(
-    server: Server,
-    session: ClientSession | undefined,
-    timeoutContext: TimeoutContext
-  ): Promise<Collection> {
     const db = this.db;
     const name = this.name;
     const options = this.options;
 
     const cmd: Document = { create: name };
-    for (const n in options) {
-      if (
-        (options as any)[n] != null &&
-        typeof (options as any)[n] !== 'function' &&
-        !ILLEGAL_COMMAND_FIELDS.has(n)
-      ) {
-        cmd[n] = (options as any)[n];
+    for (const [option, value] of Object.entries(options)) {
+      if (value != null && typeof value !== 'function' && !ILLEGAL_COMMAND_FIELDS.has(option)) {
+        cmd[option] = value;
       }
     }
+
     // otherwise just execute the command
     await super.executeCommand(server, session, cmd, timeoutContext);
     return new Collection(db, name, options);
   }
 }
 
+export async function createCollections<TSchema extends Document>(
+  db: Db,
+  name: string,
+  options: CreateCollectionOptions
+): Promise<Collection<TSchema>> {
+  const timeoutContext = TimeoutContext.create({
+    session: options.session,
+    serverSelectionTimeoutMS: db.client.s.options.serverSelectionTimeoutMS,
+    waitQueueTimeoutMS: db.client.s.options.waitQueueTimeoutMS,
+    timeoutMS: options.timeoutMS
+  });
+
+  const encryptedFields: Document | undefined =
+    options.encryptedFields ??
+    db.client.s.options.autoEncryption?.encryptedFieldsMap?.[`${db.databaseName}.${name}`];
+
+  if (encryptedFields) {
+    // Create auxilliary collections for queryable encryption support.
+    const escCollection = encryptedFields.escCollection ?? `enxcol_.${name}.esc`;
+    const ecocCollection = encryptedFields.ecocCollection ?? `enxcol_.${name}.ecoc`;
+
+    for (const collectionName of [escCollection, ecocCollection]) {
+      const createOp = new CreateCollectionOperation(db, collectionName, {
+        clusteredIndex: {
+          key: { _id: 1 },
+          unique: true
+        },
+        isEncryptedCollection: true,
+        session: options.session
+      });
+      await executeOperation(db.client, createOp, timeoutContext);
+    }
+
+    if (!options.encryptedFields) {
+      options = { ...options, encryptedFields };
+    }
+  }
+
+  const coll = await executeOperation(
+    db.client,
+    new CreateCollectionOperation(db, name, options),
+    timeoutContext
+  );
+
+  if (encryptedFields) {
+    // Create the required index for queryable encryption support.
+    const createIndexOp = CreateIndexesOperation.fromIndexSpecification(
+      db,
+      name,
+      { __safeContent__: 1 },
+      { session: options.session }
+    );
+    await executeOperation(db.client, createIndexOp, timeoutContext);
+  }
+
+  return coll as unknown as Collection<TSchema>;
+}
+
 defineAspects(CreateCollectionOperation, [Aspect.WRITE_OPERATION]);