diff options
Diffstat (limited to 'src/mongo/scripting/bson_template_evaluator.cpp')
| -rw-r--r-- | src/mongo/scripting/bson_template_evaluator.cpp | 534 |
1 files changed, 265 insertions, 269 deletions
diff --git a/src/mongo/scripting/bson_template_evaluator.cpp b/src/mongo/scripting/bson_template_evaluator.cpp index 634a478b613..0b0a243bd00 100644 --- a/src/mongo/scripting/bson_template_evaluator.cpp +++ b/src/mongo/scripting/bson_template_evaluator.cpp @@ -36,299 +36,295 @@ namespace mongo { - using std::string; - - void BsonTemplateEvaluator::initializeEvaluator() { - addOperator("RAND_INT", &BsonTemplateEvaluator::evalRandInt); - addOperator("RAND_INT_PLUS_THREAD", &BsonTemplateEvaluator::evalRandPlusThread); - addOperator("SEQ_INT", &BsonTemplateEvaluator::evalSeqInt); - addOperator("RAND_STRING", &BsonTemplateEvaluator::evalRandString); - addOperator("CONCAT", &BsonTemplateEvaluator::evalConcat); - addOperator("OID", &BsonTemplateEvaluator::evalObjId); - addOperator("VARIABLE", &BsonTemplateEvaluator::evalVariable); +using std::string; + +void BsonTemplateEvaluator::initializeEvaluator() { + addOperator("RAND_INT", &BsonTemplateEvaluator::evalRandInt); + addOperator("RAND_INT_PLUS_THREAD", &BsonTemplateEvaluator::evalRandPlusThread); + addOperator("SEQ_INT", &BsonTemplateEvaluator::evalSeqInt); + addOperator("RAND_STRING", &BsonTemplateEvaluator::evalRandString); + addOperator("CONCAT", &BsonTemplateEvaluator::evalConcat); + addOperator("OID", &BsonTemplateEvaluator::evalObjId); + addOperator("VARIABLE", &BsonTemplateEvaluator::evalVariable); +} + +BsonTemplateEvaluator::BsonTemplateEvaluator(int64_t seed) : _id(0), rng(seed) { + initializeEvaluator(); +} + +BsonTemplateEvaluator::Status BsonTemplateEvaluator::setId(size_t id) { + if (id >= (1 << 7)) + // The id is too large--doesn't fit inside 7 bits. + return StatusOpEvaluationError; + this->_id = id; + return StatusSuccess; +} + +BsonTemplateEvaluator::~BsonTemplateEvaluator() {} + +void BsonTemplateEvaluator::addOperator(const std::string& name, const OperatorFn& op) { + _operatorFunctions[name] = op; +} + +BsonTemplateEvaluator::OperatorFn BsonTemplateEvaluator::operatorEvaluator( + const std::string& op) const { + return mapFindWithDefault(_operatorFunctions, op, OperatorFn()); +} + +/* This is the top level method for using this library. It takes a BSON Object as input, + * evaluates the templates and saves the result in the builder object. + * The method returns appropriate Status on success/error condition. + */ +BsonTemplateEvaluator::Status BsonTemplateEvaluator::evaluate(const BSONObj& in, + BSONObjBuilder& builder) { + BSONForEach(e, in) { + Status st = _evalElem(e, builder); + if (st != StatusSuccess) + return st; } + return StatusSuccess; +} - BsonTemplateEvaluator::BsonTemplateEvaluator(int64_t seed) : _id(0), rng(seed) { - initializeEvaluator(); +/* Sets the variable + */ +void BsonTemplateEvaluator::setVariable(const std::string& name, const BSONElement& elem) { + this->_varMap[name] = elem.wrap(); +} + +BsonTemplateEvaluator::Status BsonTemplateEvaluator::_evalElem(const BSONElement in, + BSONObjBuilder& out) { + if (in.type() == Array) { + BSONArrayBuilder arrayBuilder(out.subarrayStart(in.fieldName())); + std::vector<BSONElement> arrElems = in.Array(); + for (unsigned int i = 0; i < arrElems.size(); i++) { + BSONElement element = arrElems[i]; + if (element.isABSONObj()) { + BSONObjBuilder newBuilder(element.objsize()); + Status st = _evalElem(element, newBuilder); + if (st != StatusSuccess) + return st; + // Only want to append the field value, not the whole object + arrayBuilder.append(newBuilder.done().getField(element.fieldName())); + } else { + arrayBuilder.append(element); + } + } + arrayBuilder.done(); + return StatusSuccess; } - BsonTemplateEvaluator::Status BsonTemplateEvaluator::setId(size_t id) { - if (id >= (1<<7)) - // The id is too large--doesn't fit inside 7 bits. - return StatusOpEvaluationError; - this->_id = id; + if (in.type() != Object) { + out.append(in); return StatusSuccess; } - BsonTemplateEvaluator::~BsonTemplateEvaluator() { } - - void BsonTemplateEvaluator::addOperator(const std::string& name, const OperatorFn& op) { - _operatorFunctions[name] = op; + // continue evaluation on subObject + BSONObj subObj = in.embeddedObject(); + + BSONObjBuilder newBuilder(subObj.objsize() + 128); + Status st = _evalObj(subObj, newBuilder); + if (st != StatusSuccess) + return st; + + BSONObj updatedObj = newBuilder.obj(); + + // check if updated object needs template evaluation + const char* opOrNot = updatedObj.firstElementFieldName(); + if (opOrNot[0] == '#') { + const char* op = opOrNot + 1; + OperatorFn fn = operatorEvaluator(op); + if (!fn) + return StatusBadOperator; + Status st = fn(this, in.fieldName(), updatedObj, out); + if (st != StatusSuccess) + return st; + } else { + out.append(in.fieldName(), updatedObj); } - - BsonTemplateEvaluator::OperatorFn BsonTemplateEvaluator::operatorEvaluator( - const std::string& op) const { - return mapFindWithDefault(_operatorFunctions, op, OperatorFn()); + return StatusSuccess; +} + +BsonTemplateEvaluator::Status BsonTemplateEvaluator::_evalObj(const BSONObj& in, + BSONObjBuilder& out) { + BSONForEach(e, in) { + Status st = _evalElem(e, out); + if (st != StatusSuccess) + return st; } - - /* This is the top level method for using this library. It takes a BSON Object as input, - * evaluates the templates and saves the result in the builder object. - * The method returns appropriate Status on success/error condition. - */ - BsonTemplateEvaluator::Status BsonTemplateEvaluator::evaluate(const BSONObj& in, - BSONObjBuilder& builder) { - BSONForEach(e, in) { - Status st = _evalElem(e, builder); - if (st != StatusSuccess) - return st; - } - return StatusSuccess; + return StatusSuccess; +} + +BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalRandInt(BsonTemplateEvaluator* btl, + const char* fieldName, + const BSONObj& in, + BSONObjBuilder& out) { + // in = { #RAND_INT: [10, 20] } + BSONObj range = in.firstElement().embeddedObject(); + if (!range["0"].isNumber() || !range["1"].isNumber()) + return StatusOpEvaluationError; + const int min = range["0"].numberInt(); + const int max = range["1"].numberInt(); + if (max <= min) + return StatusOpEvaluationError; + // range of max-min + int randomNum = min + (btl->rng.nextInt32() % (max - min)); + if (range.nFields() == 3) { + if (!range[2].isNumber()) + return StatusOpEvaluationError; + randomNum *= range[2].numberInt(); } + out.append(fieldName, randomNum); + return StatusSuccess; +} - /* Sets the variable - */ - void BsonTemplateEvaluator::setVariable(const std::string& name, const BSONElement& elem) { - this->_varMap[name] = elem.wrap(); +BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalRandPlusThread(BsonTemplateEvaluator* btl, + const char* fieldName, + const BSONObj& in, + BSONObjBuilder& out) { + // in = { #RAND_INT_PLUS_THREAD: [10, 20] } + BSONObj range = in.firstElement().embeddedObject(); + if (!range["0"].isNumber() || !range["1"].isNumber()) + return StatusOpEvaluationError; + const int min = range["0"].numberInt(); + const int max = range["1"].numberInt(); + if (max <= min) + return StatusOpEvaluationError; + int randomNum = min + (btl->rng.nextInt32() % (max - min)); + randomNum += ((max - min) * btl->_id); + out.append(fieldName, randomNum); + return StatusSuccess; +} + +BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalSeqInt(BsonTemplateEvaluator* btl, + const char* fieldName, + const BSONObj& in, + BSONObjBuilder& out) { + // in = { #SEQ_INT: { seq_id: 0, start: 10, step: -2 } + BSONObj spec = in.firstElement().embeddedObject(); + if (spec.nFields() < 3) + return StatusOpEvaluationError; + if (spec["seq_id"].eoo() || !spec["seq_id"].isNumber()) + return StatusOpEvaluationError; + if (spec["start"].eoo() || !spec["start"].isNumber()) + return StatusOpEvaluationError; + if (spec["step"].eoo() || !spec["step"].isNumber()) + return StatusOpEvaluationError; + + // If we're here, then we have a well-formed SEQ_INT specification: + // seq_id, start, and step fields, which are all numbers. + + int seq_id = spec["seq_id"].numberInt(); + long long curr_seqval = spec["start"].numberInt(); + + // Handle the optional "unique" argument. + // + // If the test requires us to keep sequences unique between different + // worker threads, then put the id number of this evaluator in the + // high order byte. + if (!spec["unique"].eoo() && spec["unique"].trueValue()) { + long long workerid = btl->_id; + curr_seqval += (workerid << ((sizeof(long long) - 1) * 8)); } - BsonTemplateEvaluator::Status BsonTemplateEvaluator::_evalElem(const BSONElement in, - BSONObjBuilder& out) { - if (in.type() == Array) { - BSONArrayBuilder arrayBuilder(out.subarrayStart(in.fieldName())); - std::vector<BSONElement> arrElems = in.Array(); - for (unsigned int i = 0; i < arrElems.size(); i++) { - BSONElement element = arrElems[i]; - if (element.isABSONObj()) { - - BSONObjBuilder newBuilder( element.objsize() ); - Status st = _evalElem(element, newBuilder); - if (st != StatusSuccess) - return st; - // Only want to append the field value, not the whole object - arrayBuilder.append(newBuilder.done().getField(element.fieldName())); - } - else { - arrayBuilder.append(element); - } - } - arrayBuilder.done(); - return StatusSuccess; - } - - if (in.type() != Object) { - out.append(in); - return StatusSuccess; - } - - //continue evaluation on subObject - BSONObj subObj = in.embeddedObject(); - - BSONObjBuilder newBuilder( subObj.objsize() + 128 ); - Status st = _evalObj(subObj, newBuilder); - if (st != StatusSuccess) - return st; - - BSONObj updatedObj = newBuilder.obj(); - - //check if updated object needs template evaluation - const char* opOrNot = updatedObj.firstElementFieldName(); - if (opOrNot[0] == '#') { - const char* op = opOrNot+1; - OperatorFn fn = operatorEvaluator(op); - if (!fn) - return StatusBadOperator; - Status st = fn(this, in.fieldName(), updatedObj, out); - if (st != StatusSuccess) - return st; - } - else { - out.append(in.fieldName(), updatedObj); - } - return StatusSuccess; - } - - BsonTemplateEvaluator::Status BsonTemplateEvaluator::_evalObj(const BSONObj& in, - BSONObjBuilder& out) { - BSONForEach(e, in) { - Status st = _evalElem(e, out); - if (st != StatusSuccess) - return st; - } - return StatusSuccess; - } - - BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalRandInt(BsonTemplateEvaluator* btl, - const char* fieldName, - const BSONObj& in, - BSONObjBuilder& out) { - // in = { #RAND_INT: [10, 20] } - BSONObj range = in.firstElement().embeddedObject(); - if (!range["0"].isNumber() || !range["1"].isNumber()) - return StatusOpEvaluationError; - const int min = range["0"].numberInt(); - const int max = range["1"].numberInt(); - if (max <= min) - return StatusOpEvaluationError; - // range of max-min - int randomNum = min + (btl->rng.nextInt32() % (max - min)); - if (range.nFields() == 3) { - if (!range[2].isNumber()) - return StatusOpEvaluationError; - randomNum *= range[2].numberInt(); - } - out.append(fieldName, randomNum); - return StatusSuccess; + if (btl->_seqIdMap.end() != btl->_seqIdMap.find(seq_id)) { + // We already have a sequence value. Add 'step' to get the next value. + int step = spec["step"].numberInt(); + curr_seqval = btl->_seqIdMap[seq_id] + step; } - BsonTemplateEvaluator::Status - BsonTemplateEvaluator::evalRandPlusThread(BsonTemplateEvaluator* btl, - const char* fieldName, - const BSONObj& in, - BSONObjBuilder& out) { - // in = { #RAND_INT_PLUS_THREAD: [10, 20] } - BSONObj range = in.firstElement().embeddedObject(); - if (!range["0"].isNumber() || !range["1"].isNumber()) - return StatusOpEvaluationError; - const int min = range["0"].numberInt(); - const int max = range["1"].numberInt(); - if (max <= min) + // Handle the optional "mod" argument. This should be done after + // handling all other options (currently just "unique"). + if (!spec["mod"].eoo()) { + if (!spec["mod"].isNumber()) return StatusOpEvaluationError; - int randomNum = min + (btl->rng.nextInt32() % (max - min)); - randomNum += ((max - min) * btl->_id); - out.append(fieldName, randomNum); - return StatusSuccess; + int modval = spec["mod"].numberInt(); + curr_seqval = (curr_seqval % modval); } - BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalSeqInt(BsonTemplateEvaluator* btl, - const char* fieldName, - const BSONObj& in, - BSONObjBuilder& out) { - // in = { #SEQ_INT: { seq_id: 0, start: 10, step: -2 } - BSONObj spec = in.firstElement().embeddedObject(); - if (spec.nFields() < 3) - return StatusOpEvaluationError; - if (spec["seq_id"].eoo() || !spec["seq_id"].isNumber()) - return StatusOpEvaluationError; - if (spec["start"].eoo() || !spec["start"].isNumber()) - return StatusOpEvaluationError; - if (spec["step"].eoo() || !spec["step"].isNumber()) - return StatusOpEvaluationError; - - // If we're here, then we have a well-formed SEQ_INT specification: - // seq_id, start, and step fields, which are all numbers. - - int seq_id = spec["seq_id"].numberInt(); - long long curr_seqval = spec["start"].numberInt(); + // Store the sequence value. + btl->_seqIdMap[seq_id] = curr_seqval; - // Handle the optional "unique" argument. - // - // If the test requires us to keep sequences unique between different - // worker threads, then put the id number of this evaluator in the - // high order byte. - if (!spec["unique"].eoo() && spec["unique"].trueValue()) { - long long workerid = btl->_id; - curr_seqval += (workerid << ((sizeof(long long) - 1) * 8)); - } - - if (btl->_seqIdMap.end() != btl->_seqIdMap.find(seq_id)) { - // We already have a sequence value. Add 'step' to get the next value. - int step = spec["step"].numberInt(); - curr_seqval = btl->_seqIdMap[seq_id] + step; - } + out.append(fieldName, curr_seqval); + return StatusSuccess; +} - // Handle the optional "mod" argument. This should be done after - // handling all other options (currently just "unique"). - if (!spec["mod"].eoo()) { - if (!spec["mod"].isNumber()) - return StatusOpEvaluationError; - int modval = spec["mod"].numberInt(); - curr_seqval = (curr_seqval % modval); - } - - // Store the sequence value. - btl->_seqIdMap[seq_id] = curr_seqval; - - out.append(fieldName, curr_seqval); - return StatusSuccess; - } - - BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalRandString(BsonTemplateEvaluator* btl, - const char* fieldName, - const BSONObj& in, - BSONObjBuilder& out) { - // in = { #RAND_STRING: [10] } - BSONObj range = in.firstElement().embeddedObject(); - if (range.nFields() != 1) - return StatusOpEvaluationError; - if (!range[0].isNumber()) - return StatusOpEvaluationError; - const int length = range["0"].numberInt(); - if (length <= 0) - return StatusOpEvaluationError; - static const char alphanum[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" - "abcdefghijklmnopqrstuvwxyz" - "0123456789+/"; - static const size_t alphaNumLength = sizeof(alphanum) - 1; - BOOST_STATIC_ASSERT(alphaNumLength == 64); - int32_t currentRand = 0; - std::string str; - for (int i = 0; i < length; ++i, currentRand >>= 6) { - if (i % 5 == 0) - currentRand = btl->rng.nextInt32(); - str.push_back(alphanum[currentRand % alphaNumLength]); - } - out.append(fieldName, str); - return StatusSuccess; - } - - BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalConcat(BsonTemplateEvaluator* btl, +BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalRandString(BsonTemplateEvaluator* btl, const char* fieldName, const BSONObj& in, BSONObjBuilder& out) { - // in = { #CONCAT: ["hello", " ", "world"] } - BSONObjBuilder objectBuilder; - Status status = btl->evaluate(in.firstElement().embeddedObject(), objectBuilder); - if (status != StatusSuccess) - return status; - BSONObj parts = objectBuilder.obj(); - if (parts.nFields() <= 1) - return StatusOpEvaluationError; - StringBuilder stringBuilder; - BSONForEach(part, parts) { - if (part.type() == String) - stringBuilder << part.String(); - else - part.toString(stringBuilder,false); - } - out.append(fieldName, stringBuilder.str()); - return StatusSuccess; + // in = { #RAND_STRING: [10] } + BSONObj range = in.firstElement().embeddedObject(); + if (range.nFields() != 1) + return StatusOpEvaluationError; + if (!range[0].isNumber()) + return StatusOpEvaluationError; + const int length = range["0"].numberInt(); + if (length <= 0) + return StatusOpEvaluationError; + static const char alphanum[] = + "ABCDEFGHIJKLMNOPQRSTUVWXYZ" + "abcdefghijklmnopqrstuvwxyz" + "0123456789+/"; + static const size_t alphaNumLength = sizeof(alphanum) - 1; + BOOST_STATIC_ASSERT(alphaNumLength == 64); + int32_t currentRand = 0; + std::string str; + for (int i = 0; i < length; ++i, currentRand >>= 6) { + if (i % 5 == 0) + currentRand = btl->rng.nextInt32(); + str.push_back(alphanum[currentRand % alphaNumLength]); } - - BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalObjId(BsonTemplateEvaluator* btl, - const char* fieldName, - const BSONObj& in, - BSONObjBuilder& out) { - // in = { #OID: 1 } - if (!mongoutils::str::equals(fieldName, "_id")) - // Error: must be generating a value for the _id field. - return StatusOpEvaluationError; - out.genOID(); - return StatusSuccess; + out.append(fieldName, str); + return StatusSuccess; +} + +BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalConcat(BsonTemplateEvaluator* btl, + const char* fieldName, + const BSONObj& in, + BSONObjBuilder& out) { + // in = { #CONCAT: ["hello", " ", "world"] } + BSONObjBuilder objectBuilder; + Status status = btl->evaluate(in.firstElement().embeddedObject(), objectBuilder); + if (status != StatusSuccess) + return status; + BSONObj parts = objectBuilder.obj(); + if (parts.nFields() <= 1) + return StatusOpEvaluationError; + StringBuilder stringBuilder; + BSONForEach(part, parts) { + if (part.type() == String) + stringBuilder << part.String(); + else + part.toString(stringBuilder, false); } - - BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalVariable(BsonTemplateEvaluator* btl, - const char* fieldName, - const BSONObj& in, - BSONObjBuilder& out) { - // in = { #VARIABLE: "x" } - BSONElement varEle = btl->_varMap[in["#VARIABLE"].String()].firstElement(); - if ( !varEle.eoo() ) { - out.appendAs(varEle, fieldName); - return StatusSuccess; - } - else { - return StatusOpEvaluationError; - } + out.append(fieldName, stringBuilder.str()); + return StatusSuccess; +} + +BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalObjId(BsonTemplateEvaluator* btl, + const char* fieldName, + const BSONObj& in, + BSONObjBuilder& out) { + // in = { #OID: 1 } + if (!mongoutils::str::equals(fieldName, "_id")) + // Error: must be generating a value for the _id field. + return StatusOpEvaluationError; + out.genOID(); + return StatusSuccess; +} + +BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalVariable(BsonTemplateEvaluator* btl, + const char* fieldName, + const BSONObj& in, + BSONObjBuilder& out) { + // in = { #VARIABLE: "x" } + BSONElement varEle = btl->_varMap[in["#VARIABLE"].String()].firstElement(); + if (!varEle.eoo()) { + out.appendAs(varEle, fieldName); + return StatusSuccess; + } else { + return StatusOpEvaluationError; } +} -} // end namespace mongo +} // end namespace mongo |