From fffd141a57b6ab09ef3ab3c17d88e3e5c7c33e61 Mon Sep 17 00:00:00 2001 From: Alyssa Clark Date: Mon, 1 Dec 2025 17:57:19 -0500 Subject: [PATCH] SERVER-109576 Implement loadable test extension for Source stages (#44059) GitOrigin-RevId: 5e484d1d058cb61f66b513e9f5b080b432333e7b --- jstests/extensions/foo_source.js | 31 ----- jstests/extensions/source.js | 105 +++++++++++++++++ src/mongo/db/extension/host/BUILD.bazel | 1 - .../db/extension/test_examples/BUILD.bazel | 2 - .../test_examples/configurations.yml | 2 +- .../extension_options/toaster.cpp | 108 +++++++++++++----- .../db/extension/test_examples/foo_source.cpp | 65 ----------- 7 files changed, 186 insertions(+), 128 deletions(-) delete mode 100644 jstests/extensions/foo_source.js create mode 100644 jstests/extensions/source.js delete mode 100644 src/mongo/db/extension/test_examples/foo_source.cpp diff --git a/jstests/extensions/foo_source.js b/jstests/extensions/foo_source.js deleted file mode 100644 index f6c22b86140..00000000000 --- a/jstests/extensions/foo_source.js +++ /dev/null @@ -1,31 +0,0 @@ -/** - * Tests that $testFooSource (source extension stage) work E2E after mongod is started with - * libfoo_source_mongo_extension.so successfully loaded. - * - * @tags: [featureFlagExtensionsAPI] - */ - -import {assertArrayEq, assertErrorCode} from "jstests/aggregation/extras/utils.js"; - -const coll = db[jsTestName()]; -coll.drop(); -const testData = [ - {_id: 0, x: 1}, - {_id: 1, x: 2}, - {_id: 2, x: 3}, -]; -assert.commandWorked(coll.insertMany(testData)); - -// Test one no-op stage passes documents through unchanged. -{ - const pipeline = [{$testFooSource: {}}]; - const result = coll.aggregate(pipeline).toArray(); - - assertArrayEq({actual: result, expected: testData}); -} - -// Test $testFoo stage fails to parse. -{ - const pipeline = [{$testFooSource: {invalidField: "value"}}]; - assertErrorCode(coll, pipeline, 11165101, "Using $testFooSource with invalid field should be rejected"); -} diff --git a/jstests/extensions/source.js b/jstests/extensions/source.js new file mode 100644 index 00000000000..1c7f19e106b --- /dev/null +++ b/jstests/extensions/source.js @@ -0,0 +1,105 @@ +/** + * Tests an extension source stage. + * + * @tags: [featureFlagExtensionsAPI] + */ + +import {FixtureHelpers} from "jstests/libs/fixture_helpers.js"; + +const collName = jsTestName(); +const coll = db[collName]; +coll.drop(); + +coll.insert({breadType: "sourdough"}); + +// Source stage must still be run against a collection. +assert.commandFailedWithCode( + db.runCommand({ + aggregate: 1, + pipeline: [{$toast: {temp: 350.0, numSlices: 2}}], + cursor: {}, + }), + ErrorCodes.InvalidNamespace, +); + +// Source stage must be first in the pipeline. +assert.commandFailedWithCode( + db.runCommand({ + aggregate: collName, + pipeline: [{$match: {breadType: "brioche"}}, {$toast: {temp: 350.0, numSlices: 2}}], + cursor: {}, + }), + 40602, +); + +// Can't have two source stages at the same level in a pipeline. +assert.commandFailedWithCode( + db.runCommand({ + aggregate: collName, + pipeline: [{$toast: {temp: 350.0, numSlices: 2}}, {$toast: {temp: 350.0, numSlices: 2}}], + cursor: {}, + }), + 40602, +); + +// EOF source stage. +let results = coll.aggregate([{$toast: {temp: 350.0, numSlices: 0}}]).toArray(); +assert.eq(results.length, 0, results); + +// Top-level source stage. +if (!FixtureHelpers.isMongos(db)) { + results = coll.aggregate([{$toast: {temp: 425.0, numSlices: 4}}]).toArray(); + assert.eq(results, [ + {slice: 0, isBurnt: true}, + {slice: 1, isBurnt: true}, + {slice: 2, isBurnt: true}, + {slice: 3, isBurnt: true}, + ]); +} else { + // This source stage will run on every shard, producing variable numbers of + // documents depending on how many shards exist in the cluster. Relax the expected + // results assertion for that case. + // TODO SERVER-114234 Remove these relaxed assertions once we can run this properly as a collectionless aggregate. + results = coll.aggregate([{$toast: {temp: 425.0, numSlices: 4}}]).toArray(); + assert.gt(results.length, 0, results); +} + +// TODO SERVER-113930 Remove failure cases and enable success cases for $lookup and $unionWith. +// Source stage in $lookup. +assert.commandFailedWithCode( + db.runCommand({ + aggregate: collName, + pipeline: [{$lookup: {from: collName, pipeline: [{$toast: {temp: 350.0, numSlices: 2}}], as: "slices"}}], + cursor: {}, + }), + 51047, +); +// results = coll.aggregate([{$lookup: {from: collName, pipeline: [{$toast: {temp: 350.0, numSlices: 2}}], as: "slices"}}]).toArray(); +// assert.eq(results, [{breadType: "sourdough", slices: [{slice: 0, isBurnt: false}, {slice: 1, isBurnt: false}]}]); + +// Source stage in $unionWith. +assert.commandFailedWithCode( + db.runCommand({ + aggregate: collName, + pipeline: [{$unionWith: {coll: collName, pipeline: [{$toast: {temp: 350.0, numSlices: 2}}]}}], + cursor: {}, + }), + 31441, +); +// results = coll.aggregate([{$unionWith: {coll: collName, pipeline: [{$toast: {temp: 350.0, numSlices: 2}}]}}]).toArray(); +// assert.eq(results, [{breadType: "sourdough"}, {slice: 0, isBurnt: false}, {slice: 1, isBurnt: false}]); + +// Source stage is not allowed in $facet. +assert.commandFailedWithCode( + db.runCommand({ + aggregate: collName, + pipeline: [{$facet: {slices: [{$toast: {temp: 250.0, numSlices: 2}}]}}], + cursor: {}, + }), + 40600, +); + +// TODO SERVER-113930 Enable this test. +// Two source stages in the pipeline. +// results = coll.aggregate([{$toast: {temp: 100.0, numSlices: 1}}, {$unionWith: {coll: collName, pipeline: [{$toast: {temp: 350.0, numSlices: 2}}]}}]).toArray(); +// assert.eq(results, [{slice: 0, notToasted: true}, {slice: 0, isBurnt: false}, {slice: 1, isBurnt: false}]); diff --git a/src/mongo/db/extension/host/BUILD.bazel b/src/mongo/db/extension/host/BUILD.bazel index 222dc2f09b9..9913515c3a9 100644 --- a/src/mongo/db/extension/host/BUILD.bazel +++ b/src/mongo/db/extension/host/BUILD.bazel @@ -81,7 +81,6 @@ mongo_cc_unit_test( "//src/mongo/db/extension/test_examples:explain_mongo_extension", "//src/mongo/db/extension/test_examples:extension_errors_mongo_extension", "//src/mongo/db/extension/test_examples:foo_mongo_extension", - "//src/mongo/db/extension/test_examples:foo_source_mongo_extension", "//src/mongo/db/extension/test_examples:host_version_fails_bad_extension", "//src/mongo/db/extension/test_examples:host_version_succeeds_mongo_extension", "//src/mongo/db/extension/test_examples:initialize_version_fails_bad_extension", diff --git a/src/mongo/db/extension/test_examples/BUILD.bazel b/src/mongo/db/extension/test_examples/BUILD.bazel index b0e3f526eaa..22b83c8d17b 100644 --- a/src/mongo/db/extension/test_examples/BUILD.bazel +++ b/src/mongo/db/extension/test_examples/BUILD.bazel @@ -42,7 +42,6 @@ extensions_with_config( ":explain_mongo_extension", ":idle_threads_mongo_extension", ":interrupt_mongo_extension", - ":foo_source_mongo_extension", ":match_topN_mongo_extension", ":native_vector_search_mongo_extension", ":metrics_mongo_extension", @@ -73,7 +72,6 @@ pkg_name = "//" + package_name() + "/" ) for extension_name in [ "foo", - "foo_source", "foo_v2", "bar", ] diff --git a/src/mongo/db/extension/test_examples/configurations.yml b/src/mongo/db/extension/test_examples/configurations.yml index a469338ad5f..2e57087c7f9 100644 --- a/src/mongo/db/extension/test_examples/configurations.yml +++ b/src/mongo/db/extension/test_examples/configurations.yml @@ -13,5 +13,5 @@ extensions: max: 10 toaster: extensionOptions: - maxToasterHeat: 5 + maxToasterHeat: 450 allowBagels: false diff --git a/src/mongo/db/extension/test_examples/extension_options/toaster.cpp b/src/mongo/db/extension/test_examples/extension_options/toaster.cpp index 5464dc5f7cb..cc093b2f2fc 100644 --- a/src/mongo/db/extension/test_examples/extension_options/toaster.cpp +++ b/src/mongo/db/extension/test_examples/extension_options/toaster.cpp @@ -28,63 +28,115 @@ */ #include "mongo/bson/bsonobj.h" +#include "mongo/db/extension/public/extension_agg_stage_static_properties_gen.h" #include "mongo/db/extension/sdk/aggregation_stage.h" #include "mongo/db/extension/sdk/extension_factory.h" #include "mongo/db/extension/sdk/test_extension_factory.h" namespace sdk = mongo::extension::sdk; -DEFAULT_LOGICAL_AST_PARSE(Toast, "$toast") -DEFAULT_LOGICAL_AST_PARSE(ToastBagel, "$toastBagel") - struct ToasterOptions { inline static double maxToasterHeat = 0; inline static bool allowBagels = false; }; -/** - * $toast is a no-op stage that requires a temperature, like {$toast: {temp: 3}}. - */ -class ToastStageDescriptor : public sdk::AggStageDescriptor { +STAGE_NAME(Toast, "$toast"); + +class ToastExecStage : public sdk::ExecAggStageSource { public: - static inline const std::string kStageName = std::string(ToastStageName); + ToastExecStage(std::string_view stageName, const mongo::BSONObj& arguments) + : sdk::ExecAggStageSource(ToastStageName), _currentSlice(0) { + _temp = arguments["temp"].Number(); + _numSlices = [&] { + if (auto numSlices = arguments["numSlices"]) { + return static_cast(numSlices.Number()); + } + return 1; + }(); + } - ToastStageDescriptor() : sdk::AggStageDescriptor(kStageName) {} + mongo::extension::ExtensionGetNextResult getNext( + const mongo::extension::sdk::QueryExecutionContextHandle& execCtx, + ::MongoExtensionExecAggStage* execStage) override { + if (_currentSlice == _numSlices) { + return mongo::extension::ExtensionGetNextResult::eof(); + } + mongo::BSONObjBuilder builder; + builder.append("slice", _currentSlice++); + if (_temp < 300.0) { + builder.append("notToasted", true); + } else { + builder.append("isBurnt", _temp > 400.0); + } + auto result = mongo::extension::ExtensionBSONObj::makeAsByteBuf(builder.obj()); + return mongo::extension::ExtensionGetNextResult::advanced(std::move(result)); + } + void open() override {} + void reopen() override {} + void close() override {} + +private: + double _temp; + int _numSlices; + int _currentSlice; +}; + +DEFAULT_LOGICAL_STAGE(Toast); + +class ToastAstNode : public sdk::TestAstNode { +public: + ToastAstNode(std::string_view stageName, const mongo::BSONObj& arguments) + : sdk::TestAstNode(stageName, arguments) {} + + // TODO SERVER-114234 Set properties for this to be a collectionless stage. + mongo::BSONObj getProperties() const override { + mongo::extension::MongoExtensionStaticProperties properties; + mongo::BSONObjBuilder builder; + properties.setPosition(mongo::extension::MongoExtensionPositionRequirementEnum::kFirst); + properties.setHostType( + mongo::extension::MongoExtensionHostTypeRequirementEnum::kRunOnceAnyNode); + properties.setRequiresInputDocSource(false); + properties.serialize(&builder); + return builder.obj(); + } +}; + +DEFAULT_PARSE_NODE(Toast); + +/** + * $toast is a source stage that requires a temperature and number of slices, like {$toast: {temp: + * 3, numSlices: 5}}. + */ +class ToastStageDescriptor : public sdk::TestStageDescriptor { +public: std::unique_ptr parse(mongo::BSONObj stageBson) const override { auto arguments = sdk::validateStageDefinition(stageBson, kStageName); - const auto obj = stageBson.getField(kStageName).Obj(); sdk_uassert(11285301, - "Failed to parse " + kStageName + ", expected {" + kStageName + - ": {temp: }}", - obj.hasField("temp") && obj.getField("temp").isNumber()); - + "expected temp input to " + kStageName, + arguments.hasField("temp") && arguments.getField("temp").isNumber()); sdk_uassert(11285302, "Failed to parse " + kStageName + ", provided temperature is higher than max " + std::to_string(ToasterOptions::maxToasterHeat), - obj.getField("temp").numberDouble() <= ToasterOptions::maxToasterHeat); + arguments["temp"].Number() <= ToasterOptions::maxToasterHeat); + if (auto numSlices = arguments["numSlices"]) { + sdk_uassert(10957601, + "numSlices must be >= 0", + numSlices.isNumber() && numSlices.Number() >= 0); + } - return std::make_unique(kStageName, arguments); + return std::make_unique(kStageName, std::move(arguments)); } }; +DEFAULT_LOGICAL_AST_PARSE(ToastBagel, "$toastBagel") + /** * $toastBagel is a no-op stage whose stage definition must be empty, like {$toastBagel: {}}. */ -class ToastBagelStageDescriptor : public sdk::AggStageDescriptor { -public: - static inline const std::string kStageName = std::string(ToastBagelStageName); - - ToastBagelStageDescriptor() : sdk::AggStageDescriptor(kStageName) {} - - std::unique_ptr parse(mongo::BSONObj stageBson) const override { - auto arguments = sdk::validateStageDefinition(stageBson, kStageName, true /* checkEmpty */); - - return std::make_unique(kStageName, arguments); - } -}; +using ToastBagelStageDescriptor = sdk::TestStageDescriptor<"$toastBagel", ToastBagelParseNode>; class ToasterExtension : public sdk::Extension { public: diff --git a/src/mongo/db/extension/test_examples/foo_source.cpp b/src/mongo/db/extension/test_examples/foo_source.cpp deleted file mode 100644 index d591a2836b3..00000000000 --- a/src/mongo/db/extension/test_examples/foo_source.cpp +++ /dev/null @@ -1,65 +0,0 @@ -/** - * Copyright (C) 2025-present MongoDB, Inc. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the Server Side Public License, version 1, - * as published by MongoDB, Inc. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * Server Side Public License for more details. - * - * You should have received a copy of the Server Side Public License - * along with this program. If not, see - * . - * - * As a special exception, the copyright holders give permission to link the - * code of portions of this program with the OpenSSL library under certain - * conditions as described in each individual source file and distribute - * linked combinations including the program with the OpenSSL library. You - * must comply with the Server Side Public License in all respects for - * all of the code used other than as permitted herein. If you modify file(s) - * with this exception, you may extend this exception to your version of the - * file(s), but you are not obligated to do so. If you do not wish to do so, - * delete this exception statement from your version. If you delete this - * exception statement from all source files in the program, then also delete - * it in the license file. - */ - -#include "mongo/bson/bsonobj.h" -#include "mongo/db/extension/sdk/aggregation_stage.h" -#include "mongo/db/extension/sdk/extension_factory.h" -#include "mongo/db/extension/sdk/test_extension_factory.h" - -namespace sdk = mongo::extension::sdk; - -DEFAULT_LOGICAL_AST_PARSE(TestFooSource, "$testFooSource") - -/** - * $testFoo is a source stage. - * - * This file is identical to foo.cpp except it is a source stage instead of a no-op stage. - */ -class TestFooSourceStageDescriptor : public sdk::AggStageDescriptor { -public: - static inline const std::string kStageName = std::string(TestFooSourceStageName); - - TestFooSourceStageDescriptor() : sdk::AggStageDescriptor(kStageName) {} - - std::unique_ptr parse(mongo::BSONObj stageBson) const override { - auto arguments = sdk::validateStageDefinition(stageBson, kStageName, true /* checkEmpty */); - - return std::make_unique(kStageName, arguments); - } -}; - -class FooExtension : public sdk::Extension { -public: - void initialize(const sdk::HostPortalHandle& portal) override { - _registerStage(portal); - } -}; - -REGISTER_EXTENSION(FooExtension) -DEFINE_GET_EXTENSION()