SERVER-112941 Add test for extensions with FLE (#44684)

GitOrigin-RevId: 35b6a203e90d725760018a374970bb402ef4db85

jstests/noPassthrough/libs/extension_helpers.js

@@ -4,6 +4,7 @@
 import {getPython3Binary} from "jstests/libs/python.js";
 import {isLinux} from "jstests/libs/os_helpers.js";
 import {ShardingTest} from "jstests/libs/shardingtest.js";
+import {ReplSetTest} from "jstests/libs/replsettest.js";
 
 export function checkPlatformCompatibleWithExtensions() {
     if (!isLinux()) {
@@ -87,7 +88,7 @@ export function deleteExtensionConfigs(names) {
  * Runs a test function in environments with one or more extension loaded, ensuring proper
  * generation and cleanup of .conf files.
  */
-export function withExtensions(extToOptionsMap, testFn) {
+export function withExtensions(extToOptionsMap, testFn, topologiesToTest = ["standalone", "sharded"]) {
     const extensionsToLoad = [];
 
     for (const [extLib, extensionOptions] of Object.entries(extToOptionsMap)) {
@@ -100,13 +101,13 @@ export function withExtensions(extToOptionsMap, testFn) {
         loadExtensions: extensionsToLoad,
     };
 
-    try {
-        {
-            const mongodConn = MongoRunner.runMongod(options);
-            testFn(mongodConn);
-            MongoRunner.stopMongod(mongodConn);
-        }
+    function runStandaloneTest(func) {
+        const mongodConn = MongoRunner.runMongod(options);
+        func(mongodConn);
+        MongoRunner.stopMongod(mongodConn);
+    }
 
+    function runShardedTest(func) {
         {
             const shardingTest = new ShardingTest({
                 shards: 1,
@@ -114,9 +115,30 @@ export function withExtensions(extToOptionsMap, testFn) {
                 config: 1,
                 mongosOptions: options,
             });
-            testFn(shardingTest.s);
+            func(shardingTest.s);
             shardingTest.stop();
         }
+    }
+
+    function runReplicaSetTest(func) {
+        const rst = new ReplSetTest({nodes: 1});
+        rst.startSet(options);
+        rst.initiate();
+        rst.awaitReplication();
+        func(rst.getPrimary());
+        rst.stopSet();
+    }
+
+    try {
+        topologiesToTest.forEach(function (topology) {
+            if (topology === "standalone") {
+                runStandaloneTest(testFn);
+            } else if (topology === "sharded") {
+                runShardedTest(testFn);
+            } else if (topology === "replica_set") {
+                runReplicaSetTest(testFn);
+            }
+        });
     } finally {
         deleteExtensionConfigs(extensionsToLoad);
     }

jstests/noPassthrough/query/fle/fle_extensions.js

@@ -0,0 +1,70 @@
+/**
+ * Tests that that extension stages are generally not supported with FLE.
+ * @tags: [featureFlagExtensionsAPI, requires_fcv_70]
+ */
+import {assertArrayEq} from "jstests/aggregation/extras/utils.js";
+import {EncryptedClient, isEnterpriseShell} from "jstests/fle2/libs/encrypted_client_util.js";
+import {checkPlatformCompatibleWithExtensions, withExtensions} from "jstests/noPassthrough/libs/extension_helpers.js";
+
+if (!isEnterpriseShell()) {
+    jsTestLog("Skipping test as it requires the enterprise module");
+    quit();
+}
+
+const aggDocs = [
+    {_id: 0, ssn: "123", name: "A", manager: "B", age: NumberLong(25), location: [0, 0]},
+    {_id: 1, ssn: "456", name: "B", manager: "C", age: NumberLong(35), location: [0, 1]},
+    {_id: 2, ssn: "789", name: "C", manager: "D", age: NumberLong(45), location: [0, 2]},
+    {_id: 3, ssn: "123", name: "D", manager: "A", age: NumberLong(55), location: [0, 3]},
+];
+
+const schema = {
+    encryptedFields: {
+        fields: [
+            {path: "ssn", bsonType: "string", queries: {queryType: "equality"}},
+            {path: "age", bsonType: "long", queries: {queryType: "equality"}},
+        ],
+    },
+};
+
+checkPlatformCompatibleWithExtensions();
+
+// Set up the encrypted collection.
+const dbName = jsTestName();
+const collName = jsTestName();
+
+function performTest(primaryConn) {
+    let db = primaryConn.getDB(dbName);
+    db.dropDatabase();
+
+    let client = new EncryptedClient(primaryConn, dbName);
+    assert.commandWorked(client.createEncryptionCollection(collName, schema));
+    let edb = client.getDB();
+
+    const encryptedColl = edb[collName];
+    for (const doc of aggDocs) {
+        assert.commandWorked(encryptedColl.einsert(doc));
+    }
+
+    // Query analysis fails with unrecognized stage error. This is expected, because query analysis
+    // does not load extensions, so it is unaware of any extensions which are not natively built
+    // into the server.
+    client.runEncryptionOperation(() => {
+        let error = assert.throws(() => encryptedColl.aggregate([{$metrics: {}}]));
+        assert.commandFailedWithCode(error, 40324);
+    });
+
+    // Verify that we can run the extension stage on the encrypted collection without
+    // auto encryption. This is a sanity check, because our expected failure results from an
+    // unrecognized stage being detected during parsing, which could happen if we failed to load the
+    // extension.
+    assertArrayEq({
+        actual: encryptedColl
+            .aggregate([{$metrics: {}}, {$project: {_id: 1}}, {$sort: {_id: 1}}, {$limit: 1}])
+            .toArray(),
+        expected: [{_id: 0}],
+    });
+}
+
+// Run tests against a replica set. Running against standalone is not supported with FLE.
+withExtensions({"libmetrics_mongo_extension.so": {}}, performTest, ["replica_set"]);