diff options
author | David Storch <david.storch@10gen.com> | 2017-03-07 14:32:07 -0500 |
---|---|---|
committer | David Storch <david.storch@10gen.com> | 2017-03-13 17:50:16 -0400 |
commit | 2f367a1e6aafa0e6bdf58f0564d3dcd7b279733f (patch) | |
tree | aed6dda9a3972b4723ea8091a5c5272a010cdff5 | |
parent | a0516b5f896703682c98cf0b8c2e333f743f4dc1 (diff) | |
download | mongo-2f367a1e6aafa0e6bdf58f0564d3dcd7b279733f.tar.gz |
SERVER-27614 add $$REMOVE agg system variable
-rw-r--r-- | buildscripts/resmokeconfig/suites/aggregation.yml | 1 | ||||
-rw-r--r-- | jstests/aggregation/variables/remove_system_variable.js | 87 | ||||
-rw-r--r-- | src/mongo/db/pipeline/document_source_add_fields_test.cpp | 42 | ||||
-rw-r--r-- | src/mongo/db/pipeline/document_source_project_test.cpp | 19 | ||||
-rw-r--r-- | src/mongo/db/pipeline/document_source_replace_root_test.cpp | 10 | ||||
-rw-r--r-- | src/mongo/db/pipeline/expression.cpp | 78 | ||||
-rw-r--r-- | src/mongo/db/pipeline/expression.h | 27 | ||||
-rw-r--r-- | src/mongo/db/pipeline/expression_test.cpp | 54 |
8 files changed, 290 insertions, 28 deletions
diff --git a/buildscripts/resmokeconfig/suites/aggregation.yml b/buildscripts/resmokeconfig/suites/aggregation.yml index c7a8c6e64be..dcfea8ecf3f 100644 --- a/buildscripts/resmokeconfig/suites/aggregation.yml +++ b/buildscripts/resmokeconfig/suites/aggregation.yml @@ -5,6 +5,7 @@ selector: - jstests/aggregation/bugs/*.js - jstests/aggregation/expressions/*.js - jstests/aggregation/sources/*/*.js + - jstests/aggregation/variables/*.js executor: js_test: diff --git a/jstests/aggregation/variables/remove_system_variable.js b/jstests/aggregation/variables/remove_system_variable.js new file mode 100644 index 00000000000..803c826af4f --- /dev/null +++ b/jstests/aggregation/variables/remove_system_variable.js @@ -0,0 +1,87 @@ +/** + * Tests for the $$REMOVE system variable. + */ +(function() { + "use strict"; + + let coll = db[jsTest.name()]; + coll.drop(); + + assert.writeOK(coll.insert({_id: 1, a: 2, b: 3})); + assert.writeOK(coll.insert({_id: 2, a: 3, b: 4})); + assert.writeOK(coll.insert({_id: 3, a: {b: 98, c: 99}})); + + let projectStage = { + $project: {_id: 0, a: 1, b: {$cond: {if: {$eq: ["$b", 4]}, then: "$$REMOVE", else: "$b"}}} + }; + + // Test that we can conditionally remove a field in $project. + assert.eq([{a: 2, b: 3}], coll.aggregate([{$match: {_id: 1}}, projectStage]).toArray()); + assert.eq([{a: 3}], coll.aggregate([{$match: {_id: 2}}, projectStage]).toArray()); + + // Test removal of a nested field, using $project. + assert.eq([{a: {b: 98}}], + coll.aggregate([{$match: {_id: 3}}, {$project: {_id: 0, "a.b": 1}}]).toArray()); + assert.eq( + [{a: {}}], + coll.aggregate([{$match: {_id: 3}}, {$project: {_id: 0, "a.b": "$$REMOVE"}}]).toArray()); + assert.eq( + [{a: {}}], + coll.aggregate([{$match: {_id: 3}}, {$project: {_id: 0, a: {b: "$$REMOVE"}}}]).toArray()); + + // Test removal of a nested field, using $addFields. + assert.eq([{_id: 3, a: {c: 99}}], + coll.aggregate([{$match: {_id: 3}}, {$addFields: {"a.b": "$$REMOVE"}}]).toArray()); + + // Test that any field path following "$$REMOVE" also evaluates to missing. + assert.eq([{_id: 3}], + coll.aggregate([{$match: {_id: 3}}, {$addFields: {"a": "$$REMOVE.a.c"}}]).toArray()); + + // Test that $$REMOVE can be used together with user-defined variables in a $let. + assert.eq([{a: {b: 3, d: 4}}], + coll.aggregate([ + {$match: {_id: 3}}, + { + $project: { + _id: 0, + a: { + $let: { + vars: {bar: 3, foo: 4}, + in : {b: "$$bar", c: "$$REMOVE", d: "$$foo"} + } + } + } + } + ]) + .toArray()); + + // Test that $$REMOVE cannot be assigned in a $let. + assert.commandFailedWithCode(db.runCommand({ + aggregate: coll.getName(), + cursor: {}, + pipeline: [ + {$match: {_id: 3}}, + {$project: {_id: 0, a: {$let: {vars: {"REMOVE": 3}, in : {b: "$$REMOVE", c: 2}}}}} + ] + }), + 16867); + + // Test that $$REMOVE, $$CURRENT, $$ROOT, and user-defined variables can all be used together. + assert.eq( + [{a: {b: 3, d: {_id: 1, a: 2, b: 3}, e: {_id: 1, a: 2, b: 3}}}], + coll.aggregate([ + {$match: {_id: 1}}, + { + $project: { + _id: 0, + a: { + $let: { + vars: {myVar: 3}, + in : {b: "$$myVar", c: "$$REMOVE", d: "$$ROOT", e: "$$CURRENT"} + } + } + } + } + ]) + .toArray()); +}()); diff --git a/src/mongo/db/pipeline/document_source_add_fields_test.cpp b/src/mongo/db/pipeline/document_source_add_fields_test.cpp index 53ff48ae906..be29dea762b 100644 --- a/src/mongo/db/pipeline/document_source_add_fields_test.cpp +++ b/src/mongo/db/pipeline/document_source_add_fields_test.cpp @@ -152,5 +152,47 @@ TEST_F(AddFieldsTest, ShouldPropagatePauses) { ASSERT_TRUE(addFields->getNext().isEOF()); } +TEST_F(AddFieldsTest, AddFieldsWithRemoveSystemVariableDoesNotAddField) { + auto addFields = DocumentSourceAddFields::create(BSON("fieldToAdd" + << "$$REMOVE"), + getExpCtx()); + auto mock = DocumentSourceMock::create(Document{{"existingField", 1}}); + addFields->setSource(mock.get()); + + auto next = addFields->getNext(); + ASSERT_TRUE(next.isAdvanced()); + Document expected{{"existingField", 1}}; + ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected); + ASSERT_TRUE(addFields->getNext().isEOF()); +} + +TEST_F(AddFieldsTest, AddFieldsWithRootSystemVariableAddsRootAsSubDoc) { + auto addFields = DocumentSourceAddFields::create(BSON("b" + << "$$ROOT"), + getExpCtx()); + auto mock = DocumentSourceMock::create(Document{{"a", 1}}); + addFields->setSource(mock.get()); + + auto next = addFields->getNext(); + ASSERT_TRUE(next.isAdvanced()); + Document expected{{"a", 1}, {"b", Document{{"a", 1}}}}; + ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected); + ASSERT_TRUE(addFields->getNext().isEOF()); +} + +TEST_F(AddFieldsTest, AddFieldsWithCurrentSystemVariableAddsRootAsSubDoc) { + auto addFields = DocumentSourceAddFields::create(BSON("b" + << "$$CURRENT"), + getExpCtx()); + auto mock = DocumentSourceMock::create(Document{{"a", 1}}); + addFields->setSource(mock.get()); + + auto next = addFields->getNext(); + ASSERT_TRUE(next.isAdvanced()); + Document expected{{"a", 1}, {"b", Document{{"a", 1}}}}; + ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected); + ASSERT_TRUE(addFields->getNext().isEOF()); +} + } // namespace } // namespace mongo diff --git a/src/mongo/db/pipeline/document_source_project_test.cpp b/src/mongo/db/pipeline/document_source_project_test.cpp index 1247de01579..60fe7a5988f 100644 --- a/src/mongo/db/pipeline/document_source_project_test.cpp +++ b/src/mongo/db/pipeline/document_source_project_test.cpp @@ -38,6 +38,7 @@ #include "mongo/db/pipeline/dependencies.h" #include "mongo/db/pipeline/document_source_mock.h" #include "mongo/db/pipeline/document_source_project.h" +#include "mongo/db/pipeline/document_value_test_util.h" #include "mongo/db/pipeline/value.h" #include "mongo/unittest/unittest.h" @@ -242,5 +243,23 @@ TEST_F(ProjectStageTest, ExclusionProjectionReportsExcludedPathsWithIdExclusion) ASSERT_EQUALS(1U, modifiedPaths.paths.count("e.f.g")); } +TEST_F(ProjectStageTest, CanUseRemoveSystemVariableToConditionallyExcludeProjectedField) { + auto project = DocumentSourceProject::create( + fromjson("{a: 1, b: {$cond: [{$eq: ['$b', 4]}, '$$REMOVE', '$b']}}"), getExpCtx()); + auto source = DocumentSourceMock::create({"{a: 2, b: 2}", "{a: 3, b: 4}"}); + project->setSource(source.get()); + auto next = project->getNext(); + ASSERT(next.isAdvanced()); + Document expected{{"a", 2}, {"b", 2}}; + ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected); + + next = project->getNext(); + ASSERT(next.isAdvanced()); + expected = Document{{"a", 3}}; + ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected); + + ASSERT(project->getNext().isEOF()); +} + } // namespace } // namespace mongo diff --git a/src/mongo/db/pipeline/document_source_replace_root_test.cpp b/src/mongo/db/pipeline/document_source_replace_root_test.cpp index ae8da121996..51411531022 100644 --- a/src/mongo/db/pipeline/document_source_replace_root_test.cpp +++ b/src/mongo/db/pipeline/document_source_replace_root_test.cpp @@ -279,6 +279,16 @@ TEST_F(ReplaceRootBasics, ReplaceRootModifiesAllFields) { ASSERT_EQUALS(0U, modifiedPaths.paths.size()); } +TEST_F(ReplaceRootBasics, ReplaceRootWithRemoveSystemVariableThrows) { + auto replaceRoot = createReplaceRoot(BSON("newRoot" + << "$$REMOVE")); + Document inputDoc = Document{{"b", 2}}; + auto mock = DocumentSourceMock::create({inputDoc}); + replaceRoot->setSource(mock.get()); + + ASSERT_THROWS_CODE(replaceRoot->getNext(), UserException, 40228); +} + /** * Fixture to test error cases of initializing the $replaceRoot stage. */ diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp index d92e4c68bd9..30cafd8ba23 100644 --- a/src/mongo/db/pipeline/expression.cpp +++ b/src/mongo/db/pipeline/expression.cpp @@ -61,9 +61,16 @@ using std::vector; /// Helper function to easily wrap constants with $const. static Value serializeConstant(Value val) { + if (val.missing()) { + return Value("$$REMOVE"_sd); + } + return Value(DOC("$const" << val)); } +const StringMap<Variables::Id> Variables::kBuiltinVarNameToId = {{"ROOT"_sd, kRootId}, + {"REMOVE"_sd, kRemoveId}}; + void Variables::uassertValidNameForUserWrite(StringData varName) { // System variables users allowed to write to (currently just one) if (varName == "CURRENT") { @@ -124,26 +131,37 @@ void Variables::uassertValidNameForUserRead(StringData varName) { } void Variables::setValue(Id id, const Value& value) { - massert(17199, "can't use Variables::setValue to set ROOT", id != ROOT_ID); - - verify(id < _numVars); - _rest[id] = value; + uassert(17199, "can't use Variables::setValue to set a reserved builtin variable", id >= 0); + auto varIndex = static_cast<size_t>(id); + invariant(varIndex < _numVars); + _rest[varIndex] = value; } Value Variables::getValue(Id id) const { - if (id == ROOT_ID) - return Value(_root); + if (id < 0) { + // This is a reserved id for a builtin variable. + switch (id) { + case Variables::kRootId: + return Value(_root); + case Variables::kRemoveId: + return Value(); + default: + MONGO_UNREACHABLE; + } + } - verify(id < _numVars); - return _rest[id]; + auto varIndex = static_cast<size_t>(id); + invariant(varIndex < _numVars); + return _rest[varIndex]; } Document Variables::getDocument(Id id) const { - if (id == ROOT_ID) + if (id == Variables::kRootId) { + // For the common case of ROOT, avoid round-tripping through Value. return _root; + } - verify(id < _numVars); - const Value var = _rest[id]; + const Value var = getValue(id); if (var.getType() == Object) return var.getDocument(); @@ -151,24 +169,34 @@ Document Variables::getDocument(Id id) const { } Variables::Id VariablesParseState::defineVariable(StringData name) { - // caller should have validated before hand by using Variables::uassertValidNameForUserWrite - massert(17275, "Can't redefine ROOT", name != "ROOT"); + // Caller should have validated before hand by using Variables::uassertValidNameForUserWrite. + massert(17275, + "Can't redefine a non-user-writable variable", + Variables::kBuiltinVarNameToId.find(name) == Variables::kBuiltinVarNameToId.end()); Variables::Id id = _idGenerator->generateId(); + invariant(id >= 0); _variables[name] = id; return id; } Variables::Id VariablesParseState::getVariable(StringData name) const { - StringMap<Variables::Id>::const_iterator it = _variables.find(name); - if (it != _variables.end()) + auto it = _variables.find(name); + if (it != _variables.end()) { + // Found a user-defined variable. return it->second; + } - uassert(17276, - str::stream() << "Use of undefined variable: " << name, - name == "ROOT" || name == "CURRENT"); + it = Variables::kBuiltinVarNameToId.find(name); + if (it != Variables::kBuiltinVarNameToId.end()) { + // This is a builtin variable. + return it->second; + } - return Variables::ROOT_ID; + // If we didn't find either a user-defined or builtin variable, then we reject everything other + // than CURRENT. If this is CURRENT, then we treat it as equivalent to ROOT. + uassert(17276, str::stream() << "Use of undefined variable: " << name, name == "CURRENT"); + return Variables::kRootId; } /* --------------------------- Expression ------------------------------ */ @@ -1419,7 +1447,7 @@ Value ExpressionObject::serialize(bool explain) const { // this is the old deprecated version only used by tests not using variables intrusive_ptr<ExpressionFieldPath> ExpressionFieldPath::create( const boost::intrusive_ptr<ExpressionContext>& expCtx, const string& fieldPath) { - return new ExpressionFieldPath(expCtx, "CURRENT." + fieldPath, Variables::ROOT_ID); + return new ExpressionFieldPath(expCtx, "CURRENT." + fieldPath, Variables::kRootId); } // this is the new version that supports every syntax @@ -1454,12 +1482,16 @@ ExpressionFieldPath::ExpressionFieldPath(const boost::intrusive_ptr<ExpressionCo : Expression(expCtx), _fieldPath(theFieldPath), _variable(variable) {} intrusive_ptr<Expression> ExpressionFieldPath::optimize() { - /* nothing can be done for these */ + if (_variable == Variables::kRemoveId) { + // The REMOVE system variable optimizes to a constant missing value. + return ExpressionConstant::create(getExpressionContext(), Value()); + } + return intrusive_ptr<Expression>(this); } void ExpressionFieldPath::addDependencies(DepsTracker* deps) const { - if (_variable == Variables::ROOT_ID) { // includes CURRENT when it is equivalent to ROOT. + if (_variable == Variables::kRootId) { // includes CURRENT when it is equivalent to ROOT. if (_fieldPath.getPathLength() == 1) { deps->needWholeDocument = true; // need full doc if just "$$ROOT" } else { @@ -1511,7 +1543,7 @@ Value ExpressionFieldPath::evaluateInternal(Variables* vars) const { if (_fieldPath.getPathLength() == 1) // get the whole variable return vars->getValue(_variable); - if (_variable == Variables::ROOT_ID) { + if (_variable == Variables::kRootId) { // ROOT is always a document so use optimized code path return evaluatePath(1, vars->getRoot()); } diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h index 0e8185ef3db..e876c5e7079 100644 --- a/src/mongo/db/pipeline/expression.h +++ b/src/mongo/db/pipeline/expression.h @@ -72,11 +72,15 @@ class Variables { public: /** - * Each unique variable is assigned a unique id of this type + * Each unique variable is assigned a unique id of this type. Negative ids are reserved for + * system variables and non-negative ids are allocated for user variables. */ - typedef size_t Id; + typedef int64_t Id; - // This is only for expressions that use no variables (even ROOT). + /** + * Constructs a placeholder for expressions that use no variables (even builtins like ROOT or + * REMOVE). + */ Variables() : _numVars(0) {} explicit Variables(size_t numVars, const Document& root = Document()) @@ -85,7 +89,12 @@ public: static void uassertValidNameForUserWrite(StringData varName); static void uassertValidNameForUserRead(StringData varName); - static const Id ROOT_ID = Id(-1); + // Ids for builtin variables. + static constexpr Id kRootId = Id(-1); + static constexpr Id kRemoveId = Id(-2); + + // Map from builtin var name to reserved id number. + static const StringMap<Id> kBuiltinVarNameToId; /** * Use this instead of setValue for setting ROOT @@ -100,11 +109,19 @@ public: return _root; } + /** + * Sets the value of a user-defined variable. Illegal to use with the reserved builtin variables + * defined above. + */ void setValue(Id id, const Value& value); + + /** + * Gets the value of a user-defined or system variable. + */ Value getValue(Id id) const; /** - * returns Document() for non-document values. + * Returns Document() for non-document values, but otherwise identical to getValue(). */ Document getDocument(Id id) const; diff --git a/src/mongo/db/pipeline/expression_test.cpp b/src/mongo/db/pipeline/expression_test.cpp index 315a2faeac7..f08617396bb 100644 --- a/src/mongo/db/pipeline/expression_test.cpp +++ b/src/mongo/db/pipeline/expression_test.cpp @@ -2074,6 +2074,18 @@ private: } }; +TEST(ExpressionConstantTest, ConstantOfValueMissingRemovesField) { + intrusive_ptr<Expression> expression = ExpressionConstant::create(nullptr, Value()); + assertBinaryEqual(BSONObj(), toBson(expression->evaluate(Document{{"foo", Value("bar"_sd)}}))); +} + +TEST(ExpressionConstantTest, ConstantOfValueMissingSerializesToRemoveSystemVar) { + intrusive_ptr<Expression> expression = ExpressionConstant::create(nullptr, Value()); + assertBinaryEqual(BSON("field" + << "$$REMOVE"), + BSON("field" << expression->serialize(false))); +} + } // namespace Constant TEST(ExpressionFromAccumulators, Avg) { @@ -3836,6 +3848,48 @@ TEST(ExpressionTypeTest, WithMaxKeyValue) { } // namespace Type +namespace BuiltinRemoveVariable { + +TEST(BuiltinRemoveVariableTest, TypeOfRemoveIsMissing) { + assertExpectedResults("$type", {{{Value("$$REMOVE"_sd)}, Value("missing"_sd)}}); +} + +TEST(BuiltinRemoveVariableTest, LiteralEscapesRemoveVar) { + assertExpectedResults( + "$literal", {{{Value("$$REMOVE"_sd)}, Value(std::vector<Value>{Value("$$REMOVE"_sd)})}}); +} + +TEST(BuiltinRemoveVariableTest, RemoveSerializesCorrectly) { + VariablesIdGenerator idGenerator{}; + VariablesParseState vps{&idGenerator}; + auto expression = ExpressionFieldPath::parse(nullptr, "$$REMOVE", vps); + ASSERT_BSONOBJ_EQ(BSON("foo" + << "$$REMOVE"), + BSON("foo" << expression->serialize(false))); +} + +TEST(BuiltinRemoveVariableTest, RemoveSerializesCorrectlyWithTrailingPath) { + VariablesIdGenerator idGenerator{}; + VariablesParseState vps{&idGenerator}; + auto expression = ExpressionFieldPath::parse(nullptr, "$$REMOVE.a.b", vps); + ASSERT_BSONOBJ_EQ(BSON("foo" + << "$$REMOVE.a.b"), + BSON("foo" << expression->serialize(false))); +} + +TEST(BuiltinRemoveVariableTest, RemoveSerializesCorrectlyAfterOptimization) { + VariablesIdGenerator idGenerator{}; + VariablesParseState vps{&idGenerator}; + auto expression = ExpressionFieldPath::parse(nullptr, "$$REMOVE.a.b", vps); + auto optimizedExpression = expression->optimize(); + ASSERT(dynamic_cast<ExpressionConstant*>(optimizedExpression.get())); + ASSERT_BSONOBJ_EQ(BSON("foo" + << "$$REMOVE"), + BSON("foo" << optimizedExpression->serialize(false))); +} + +} // namespace BuiltinRemoveVariable + namespace ToLower { class ExpectedResultBase { |