diff options
Diffstat (limited to 'src/mongo/s/write_ops')
-rw-r--r-- | src/mongo/s/write_ops/batched_command_request.cpp | 45 | ||||
-rw-r--r-- | src/mongo/s/write_ops/batched_command_request.h | 10 | ||||
-rw-r--r-- | src/mongo/s/write_ops/batched_insert_request.cpp | 6 | ||||
-rw-r--r-- | src/mongo/s/write_ops/batched_insert_request.h | 1 | ||||
-rw-r--r-- | src/mongo/s/write_ops/batched_insert_request_test.cpp | 88 |
5 files changed, 142 insertions, 8 deletions
diff --git a/src/mongo/s/write_ops/batched_command_request.cpp b/src/mongo/s/write_ops/batched_command_request.cpp index 92f6f99069f..8b976c851c7 100644 --- a/src/mongo/s/write_ops/batched_command_request.cpp +++ b/src/mongo/s/write_ops/batched_command_request.cpp @@ -275,6 +275,51 @@ namespace mongo { INVOKE( getMetadata ); } + /** + * Generates a new request with insert _ids if required. Otherwise returns NULL. + */ + BatchedCommandRequest* // + BatchedCommandRequest::cloneWithIds(const BatchedCommandRequest& origCmdRequest) { + + if (origCmdRequest.getBatchType() != BatchedCommandRequest::BatchType_Insert + || origCmdRequest.isInsertIndexRequest()) + return NULL; + + auto_ptr<BatchedInsertRequest> idRequest; + BatchedInsertRequest* origRequest = origCmdRequest.getInsertRequest(); + + const vector<BSONObj>& inserts = origRequest->getDocuments(); + + size_t i = 0u; + for (vector<BSONObj>::const_iterator it = inserts.begin(); it != inserts.end(); ++it, ++i) { + + const BSONObj& insert = *it; + BSONObj idInsert; + + if (insert["_id"].eoo()) { + BSONObjBuilder idInsertB; + idInsertB.append("_id", OID::gen()); + idInsertB.appendElements(insert); + idInsert = idInsertB.obj(); + } + + if (NULL == idRequest.get() && !idInsert.isEmpty()) { + idRequest.reset(new BatchedInsertRequest); + origRequest->cloneTo(idRequest.get()); + } + + if (!idInsert.isEmpty()) { + idRequest->setDocumentAt(i, idInsert); + } + } + + if (NULL == idRequest.get()) + return NULL; + + // Command request owns idRequest + return new BatchedCommandRequest(idRequest.release()); + } + bool BatchedCommandRequest::containsUpserts( const BSONObj& writeCmdObj ) { BSONElement updatesEl = writeCmdObj[BatchedUpdateRequest::updates()]; diff --git a/src/mongo/s/write_ops/batched_command_request.h b/src/mongo/s/write_ops/batched_command_request.h index 7b59268484f..a78b13d1197 100644 --- a/src/mongo/s/write_ops/batched_command_request.h +++ b/src/mongo/s/write_ops/batched_command_request.h @@ -140,6 +140,16 @@ namespace mongo { BatchedRequestMetadata* getMetadata() const; // + // Helpers for batch pre-processing + // + + /** + * Generates a new request, the same as the old, but with insert _ids if required. + * Returns NULL if this is not an insert request or all inserts already have _ids. + */ + static BatchedCommandRequest* cloneWithIds(const BatchedCommandRequest& origCmdRequest); + + // // Helpers for auth pre-parsing // diff --git a/src/mongo/s/write_ops/batched_insert_request.cpp b/src/mongo/s/write_ops/batched_insert_request.cpp index 75da1642db1..f2e47ca0c2b 100644 --- a/src/mongo/s/write_ops/batched_insert_request.cpp +++ b/src/mongo/s/write_ops/batched_insert_request.cpp @@ -239,6 +239,12 @@ namespace mongo { return _documents.at(pos); } + void BatchedInsertRequest::setDocumentAt(size_t pos, const BSONObj& doc) { + dassert(_isDocumentsSet); + dassert(_documents.size() > pos); + _documents[pos] = doc; + } + void BatchedInsertRequest::setWriteConcern(const BSONObj& writeConcern) { _writeConcern = writeConcern.getOwned(); _isWriteConcernSet = true; diff --git a/src/mongo/s/write_ops/batched_insert_request.h b/src/mongo/s/write_ops/batched_insert_request.h index 8fe040552d0..b26f05d9236 100644 --- a/src/mongo/s/write_ops/batched_insert_request.h +++ b/src/mongo/s/write_ops/batched_insert_request.h @@ -97,6 +97,7 @@ namespace mongo { std::size_t sizeDocuments() const; const std::vector<BSONObj>& getDocuments() const; const BSONObj& getDocumentsAt(std::size_t pos) const; + void setDocumentAt(std::size_t pos, const BSONObj& doc); void setWriteConcern(const BSONObj& writeConcern); void unsetWriteConcern(); diff --git a/src/mongo/s/write_ops/batched_insert_request_test.cpp b/src/mongo/s/write_ops/batched_insert_request_test.cpp index 20ff1319f23..b326cfcfed9 100644 --- a/src/mongo/s/write_ops/batched_insert_request_test.cpp +++ b/src/mongo/s/write_ops/batched_insert_request_test.cpp @@ -32,19 +32,14 @@ #include "mongo/db/jsobj.h" #include "mongo/unittest/unittest.h" +#include "mongo/s/write_ops/batched_command_request.h" namespace { - using mongo::BSONArray; - using mongo::BSONObj; - using mongo::BatchedInsertRequest; - using mongo::BatchedRequestMetadata; - using mongo::BSONArrayBuilder; - using mongo::OID; - using mongo::OpTime; + using namespace mongo; + using std::auto_ptr; using std::string; - TEST(RoundTrip, Normal) { BSONArray insertArray = BSON_ARRAY(BSON("a" << 1) << BSON("b" << 1)); @@ -75,4 +70,81 @@ namespace { ASSERT_EQUALS(0, genInsertRequestObj.woCompare(origInsertRequestObj)); } + TEST(GenID, All) { + + BatchedCommandRequest cmdRequest(BatchedCommandRequest::BatchType_Insert); + BatchedInsertRequest& request = *cmdRequest.getInsertRequest(); + + request.setCollName("foo.bar"); + request.setOrdered(false); + + BSONObj insertA = BSON( "a" << 1 ); + BSONObj insertB = BSON( "b" << 1 ); + request.addToDocuments(insertA); + request.addToDocuments(insertB); + + auto_ptr<BatchedCommandRequest> idCmdRequest; + idCmdRequest.reset(BatchedCommandRequest::cloneWithIds(cmdRequest)); + ASSERT(idCmdRequest.get()); + + BatchedInsertRequest* idRequest = idCmdRequest->getInsertRequest(); + ASSERT_EQUALS(idRequest->getCollName(), request.getCollName()); + ASSERT_EQUALS(idRequest->getOrdered(), request.getOrdered()); + + ASSERT(!idRequest->getDocumentsAt(0)["_id"].eoo()); + ASSERT_EQUALS(idRequest->getDocumentsAt(0).nFields(), 2); + ASSERT(!idRequest->getDocumentsAt(1)["_id"].eoo()); + ASSERT_EQUALS(idRequest->getDocumentsAt(1).nFields(), 2); + } + + TEST(GenID, Partial) { + + BatchedCommandRequest cmdRequest(BatchedCommandRequest::BatchType_Insert); + BatchedInsertRequest& request = *cmdRequest.getInsertRequest(); + + request.setCollName("foo.bar"); + request.setOrdered(false); + + BSONObj insertA = BSON( "a" << 1 ); + BSONObj insertB = BSON( "b" << 1 << "_id" << 1 ); + BSONObj insertC = BSON( "c" << 1 ); + request.addToDocuments(insertA); + request.addToDocuments(insertB); + request.addToDocuments(insertC); + + auto_ptr<BatchedCommandRequest> idCmdRequest; + idCmdRequest.reset(BatchedCommandRequest::cloneWithIds(cmdRequest)); + ASSERT(idCmdRequest.get()); + + BatchedInsertRequest* idRequest = idCmdRequest->getInsertRequest(); + ASSERT_EQUALS(idRequest->getCollName(), request.getCollName()); + ASSERT_EQUALS(idRequest->getOrdered(), request.getOrdered()); + + ASSERT(!idRequest->getDocumentsAt(0)["_id"].eoo()); + ASSERT_EQUALS(idRequest->getDocumentsAt(0).nFields(), 2); + ASSERT(!idRequest->getDocumentsAt(1)["_id"].eoo()); + ASSERT_EQUALS(idRequest->getDocumentsAt(1).nFields(), 2); + ASSERT(!idRequest->getDocumentsAt(2)["_id"].eoo()); + ASSERT_EQUALS(idRequest->getDocumentsAt(1).nFields(), 2); + } + + TEST(GenID, None) { + + BatchedCommandRequest cmdRequest(BatchedCommandRequest::BatchType_Insert); + BatchedInsertRequest& request = *cmdRequest.getInsertRequest(); + + // We need to check for system.indexes namespace + request.setCollName("foo.bar"); + + BSONObj insertA = BSON( "_id" << 0 << "a" << 1 ); + BSONObj insertB = BSON( "b" << 1 << "_id" << 1 ); + request.addToDocuments(insertA); + request.addToDocuments(insertB); + + auto_ptr<BatchedCommandRequest> idCmdRequest; + idCmdRequest.reset(BatchedCommandRequest::cloneWithIds(cmdRequest)); + ASSERT(!idCmdRequest.get()); + } + + } // unnamed namespace |