SERVER-13518 generate _ids in write batches if not set in mongos

(cherry picked from commit 4b589d8aab1675e03dc7efafc32239bda660fe96)

6 files changed, 147 insertions, 9 deletions

diff --git a/src/mongo/s/cluster_write.cpp b/src/mongo/s/cluster_write.cpp
index 8b2008f02b8..15cd83c2be0 100644
--- a/src/mongo/s/cluster_write.cpp
+++ b/src/mongo/s/cluster_write.cpp
@@ -315,9 +315,13 @@ namespace mongo {
         dassert( response->isValid(NULL) );
     }
 
-    void ClusterWriter::write( const BatchedCommandRequest& request,
+    void ClusterWriter::write( const BatchedCommandRequest& origRequest,
                                BatchedCommandResponse* response ) {
 
+        // Add _ids to insert request if req'd
+        auto_ptr<BatchedCommandRequest> idRequest(BatchedCommandRequest::cloneWithIds(origRequest));
+        const BatchedCommandRequest& request = NULL != idRequest.get() ? *idRequest : origRequest;
+
         const NamespaceString nss = NamespaceString( request.getNS() );
         if ( !nss.isValid() ) {
             toBatchError( Status( ErrorCodes::InvalidNamespace,
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