SERVER-112826 Implement e2e js tests for desugar extension stage query shape (#43718)

GitOrigin-RevId: be189eef31d17c12faeda4c4284692a9cf5ba868
This commit is contained in:
Finley Lau 2025-11-09 10:06:42 -06:00 committed by MongoDB Bot
parent e537c08bbf
commit 6651933d77
2 changed files with 66 additions and 0 deletions

View File

@ -18,6 +18,14 @@ function insertEmptyShape(coll) {
return [{"$shapify": {}}];
}
function insertShapeForDesugarStage(coll) {
const spec = {$shapifyDesugar: {a: 1, b: 2}};
coll.aggregate([spec]).toArray();
// Expected shape.
return [spec];
}
function insertLiteralShape(coll) {
// Shape with literals. Test different values to ensure they all map to the same shape.
let literalShape = {
@ -245,3 +253,6 @@ statsSoFar = insertShapeAndValidateResults(testDb, insertShapeWithIdentifiers3,
statsSoFar = insertShapeAndValidateResults(testDb, insertShapeWithTransformedIdentifiers4, statsSoFar, true);
statsSoFar = insertShapeAndValidateResults(testDb, insertShapeWithTransformedIdentifiers5, statsSoFar, true);
statsSoFar = insertShapeAndValidateResults(testDb, insertShapeWithTransformedIdentifiers6, statsSoFar, true);
// Test that a desugar stage's query shape is calculated from the pre-desugared stage, not the post-desugared stages.
statsSoFar = insertShapeAndValidateResults(testDb, insertShapeForDesugarStage, statsSoFar);

View File

@ -50,6 +50,15 @@ using namespace mongo;
*/
static constexpr std::string kShapifyStageName = "$shapify";
/**
* Extension desugar stage with a non-default query shape implementation. Syntax:
*
* {$shapifyDesugar: {random object>}}
*
* ShapifyDesugar will expand into multiple $shapify stages for query shape testing purposes.
*/
static constexpr std::string kShapifyDesugarStageName = "$shapifyDesugar";
class ShapifyLogicalStage : public sdk::LogicalAggStage {
public:
ShapifyLogicalStage(BSONObj input) : _input(input) {}
@ -153,10 +162,56 @@ public:
}
};
class ShapifyDesugarParseNode : public sdk::AggStageParseNode {
public:
ShapifyDesugarParseNode(BSONObj input)
: sdk::AggStageParseNode(kShapifyDesugarStageName), _input(input) {}
size_t getExpandedSize() const override {
return 3;
}
std::vector<mongo::extension::VariantNodeHandle> expand() const override {
std::vector<mongo::extension::VariantNodeHandle> expanded;
expanded.reserve(getExpandedSize());
BSONObj shapifySpec =
BSON(kShapifyStageName << BSON("int" << 1 << "ident_name" << "Alice"));
expanded.emplace_back(
new sdk::ExtensionAggStageParseNode(std::make_unique<ShapifyParseNode>(shapifySpec)));
expanded.emplace_back(
new sdk::ExtensionAggStageParseNode(std::make_unique<ShapifyParseNode>(shapifySpec)));
expanded.emplace_back(
new sdk::ExtensionAggStageParseNode(std::make_unique<ShapifyParseNode>(shapifySpec)));
return expanded;
}
BSONObj getQueryShape(const ::MongoExtensionHostQueryShapeOpts* ctx) const override {
return BSON(kShapifyDesugarStageName << _input);
}
private:
BSONObj _input;
};
class ShapifyDesugarStageDescriptor : public sdk::AggStageDescriptor {
public:
static inline const std::string kStageName = std::string(kShapifyDesugarStageName);
ShapifyDesugarStageDescriptor() : sdk::AggStageDescriptor(kStageName) {}
std::unique_ptr<sdk::AggStageParseNode> parse(mongo::BSONObj stageBson) const override {
sdk::validateStageDefinition(stageBson, kStageName);
return std::make_unique<ShapifyDesugarParseNode>(stageBson[kStageName].Obj().getOwned());
}
};
class ShapifyExtension : public sdk::Extension {
public:
void initialize(const sdk::HostPortalHandle& portal) override {
_registerStage<ShapifyStageDescriptor>(portal);
_registerStage<ShapifyDesugarStageDescriptor>(portal);
}
};