From 2926a9228a4f4bf5ec4ecc9e61d3a9af7ef4f022 Mon Sep 17 00:00:00 2001
From: bailey <bailey.pearson@mongodb.com>
Date: Fri, 24 Jan 2025 15:42:15 -0700
Subject: [PATCH] add discriminator support

---
 lib/drivers/node-mongodb-native/connection.js | 28 +++++++++++++------
 test/encryption/encryption.test.js            | 22 +++++++--------
 2 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/lib/drivers/node-mongodb-native/connection.js b/lib/drivers/node-mongodb-native/connection.js
index 0e6515a321..a27ee3f52f 100644
--- a/lib/drivers/node-mongodb-native/connection.js
+++ b/lib/drivers/node-mongodb-native/connection.js
@@ -12,6 +12,7 @@ const pkg = require('../../../package.json');
 const processConnectionOptions = require('../../helpers/processConnectionOptions');
 const setTimeout = require('../../helpers/timers').setTimeout;
 const utils = require('../../utils');
+const { Schema } = require('../../mongoose');
 
 /**
  * A [node-mongodb-native](https://github.com/mongodb/node-mongodb-native) connection implementation.
@@ -357,20 +358,31 @@ NativeConnection.prototype.createClient = async function createClient(uri, optio
  * options.
   */
 NativeConnection.prototype._buildEncryptionSchemas = function() {
-  const schemaMap = Object.values(this.models).filter(model => model.schema.encryptionType() === 'csfle').reduce(
-    (schemaMap, model) => {
-      const { schema, collection: { collectionName } } = model;
-      const namespace = `${this.$dbName}.${collectionName}`;
+  const qeMappings = {};
+  const csfleMappings = {};
+
+  for (const model of Object.values(this.models)) {
+    const { schema, collection: { collectionName } } = model;
+    const namespace = `${this.$dbName}.${collectionName}`;
+    if (schema.encryptionType() === 'csfle') {
+      csfleMappings[namespace] ??= new Schema({}, { encryptionType: 'csfle' });
+      csfleMappings[namespace].add(schema);
+    } else if (schema.encryptionType() === 'queryableEncryption') {
+      qeMappings[namespace] ??= new Schema({}, { encryptionType: 'queryableEncryption' });
+      qeMappings[namespace].add(schema);
+    }
+  }
+
+  const schemaMap = Object.entries(csfleMappings).reduce(
+    (schemaMap, [namespace, schema]) => {
       schemaMap[namespace] = schema._buildSchemaMap();
       return schemaMap;
     },
     {}
   );
 
-  const encryptedFieldsMap = Object.values(this.models).filter(model => model.schema.encryptionType() === 'qe').reduce(
-    (encryptedFieldsMap, model) => {
-      const { schema, collection: { collectionName } } = model;
-      const namespace = `${this.$dbName}.${collectionName}`;
+  const encryptedFieldsMap = Object.entries(qeMappings).reduce(
+    (encryptedFieldsMap, [namespace, schema]) => {
       encryptedFieldsMap[namespace] = schema._buildEncryptedFields();
       return encryptedFieldsMap;
     },
diff --git a/test/encryption/encryption.test.js b/test/encryption/encryption.test.js
index ea20a9248d..4efe48fae6 100644
--- a/test/encryption/encryption.test.js
+++ b/test/encryption/encryption.test.js
@@ -124,14 +124,14 @@ describe('ci', () => {
         it(`${name} encrypts and decrypts`, test);
       });
 
-      describe('QE', function() {
+      describe('queryableEncryption', function() {
         beforeEach(async function() {
           schema = new Schema({
             field: {
               type, encrypt: { keyId: keyId }
             }
           }, {
-            encryptionType: 'qe'
+            encryptionType: 'queryableEncryption'
           });
 
           connection = createConnection();
@@ -187,7 +187,7 @@ describe('ci', () => {
                 }
               }
             }, {
-              encryptionType: 'qe'
+              encryptionType: 'queryableEncryption'
             });
 
             connection = createConnection();
@@ -231,12 +231,12 @@ describe('ci', () => {
                 }
               }
             }, {
-              encryptionType: 'qe'
+              encryptionType: 'queryableEncryption'
             });
             const schema = new Schema({
               a: nestedSchema
             }, {
-              encryptionType: 'qe'
+              encryptionType: 'queryableEncryption'
             });
 
             connection = createConnection();
@@ -307,7 +307,7 @@ describe('ci', () => {
                 }
               }
             }, {
-              encryptionType: 'qe'
+              encryptionType: 'queryableEncryption'
             });
 
             connection = createConnection();
@@ -398,7 +398,7 @@ describe('ci', () => {
                 }
               }
             }, {
-              encryptionType: 'qe'
+              encryptionType: 'queryableEncryption'
             });
 
             connection = createConnection();
@@ -481,7 +481,7 @@ describe('ci', () => {
                 }
               }
             }, {
-              encryptionType: 'qe'
+              encryptionType: 'queryableEncryption'
             }));
             const model2 = connection.model('Model2', new Schema({
               b: {
@@ -491,7 +491,7 @@ describe('ci', () => {
                 }
               }
             }, {
-              encryptionType: 'qe'
+              encryptionType: 'queryableEncryption'
             }));
 
             return { model1, model2 };
@@ -551,7 +551,7 @@ describe('ci', () => {
             }
           }
         }, {
-          encryptionType: 'qe'
+          encryptionType: 'queryableEncryption'
         }));
         const model2 = connection.model('Model2', new Schema({
           b: {
@@ -684,7 +684,7 @@ describe('ci', () => {
       });
 
 
-      describe('qe', function() {
+      describe('queryableEncryption', function() {
         beforeEach(async function() {
           connection = createConnection();