SERVER-29254 add findSingleton and putSingleton methods to StorageInterface

5 files changed, 219 insertions, 0 deletions

diff --git a/src/mongo/db/repl/storage_interface.h b/src/mongo/db/repl/storage_interface.h
index 12f0d25507f..bf5a101b6a7 100644
--- a/src/mongo/db/repl/storage_interface.h
+++ b/src/mongo/db/repl/storage_interface.h
@@ -270,6 +270,27 @@ public:
                                                              std::size_t limit) = 0;
 
     /**
+     * Finds a singleton document in a collection. Returns 'CollectionIsEmpty' if the collection
+     * is empty or 'TooManyMatchingDocuments' if it is not a singleton collection.
+     */
+    virtual StatusWith<BSONObj> findSingleton(OperationContext* opCtx,
+                                              const NamespaceString& nss) = 0;
+
+    /**
+     * Updates a singleton document in a collection. Upserts the document if it does not exist. If
+     * the document is upserted and no '_id' is provided, one will be generated.
+     * If the collection has more than 1 document, the update will only be performed on the first
+     * one found.
+     * Returns 'NamespaceNotFound' if the collection does not exist. This does not implicitly
+     * create the collection so that the caller can create the collection with any collection
+     * options they want (ex: capped, temp, collation, etc.).
+     */
+    virtual Status putSingleton(OperationContext* opCtx,
+                                const NamespaceString& nss,
+                                const BSONObj& update) = 0;
+
+
+    /**
      * Finds a single document in the collection referenced by the specified _id.
      *
      * Not supported on collections with a default collation.
diff --git a/src/mongo/db/repl/storage_interface_impl.cpp b/src/mongo/db/repl/storage_interface_impl.cpp
index 0144fb76cb4..4ab7100f2b3 100644
--- a/src/mongo/db/repl/storage_interface_impl.cpp
+++ b/src/mongo/db/repl/storage_interface_impl.cpp
@@ -729,6 +729,31 @@ StatusWith<std::vector<BSONObj>> StorageInterfaceImpl::deleteDocuments(
                                   FindDeleteMode::kDelete);
 }
 
+StatusWith<BSONObj> StorageInterfaceImpl::findSingleton(OperationContext* opCtx,
+                                                        const NamespaceString& nss) {
+    auto result = findDocuments(opCtx,
+                                nss,
+                                boost::none,  // Collection scan.
+                                StorageInterface::ScanDirection::kForward,
+                                {},  // Start at the beginning of the collection.
+                                BoundInclusion::kIncludeStartKeyOnly,
+                                2U);  // Ask for 2 documents to ensure it's a singleton.
+    if (!result.isOK()) {
+        return result.getStatus();
+    }
+
+    const auto& docs = result.getValue();
+    if (docs.empty()) {
+        return {ErrorCodes::CollectionIsEmpty,
+                str::stream() << "No document found in namespace: " << nss.ns()};
+    } else if (docs.size() != 1U) {
+        return {ErrorCodes::TooManyMatchingDocuments,
+                str::stream() << "More than singleton document found in namespace: " << nss.ns()};
+    }
+
+    return docs.front();
+}
+
 StatusWith<BSONObj> StorageInterfaceImpl::findById(OperationContext* opCtx,
                                                    const NamespaceString& nss,
                                                    const BSONElement& idKey) {
@@ -760,6 +785,53 @@ StatusWith<BSONObj> makeUpsertQuery(const BSONElement& idKey) {
     return query;
 }
 
+Status _upsertWithQuery(OperationContext* opCtx,
+                        const NamespaceString& nss,
+                        const BSONObj& query,
+                        const BSONObj& update) {
+    UpdateRequest request(nss);
+    request.setQuery(query);
+    request.setUpdates(update);
+    request.setUpsert(true);
+    invariant(!request.isMulti());  // We only want to update one document for performance.
+    invariant(!request.shouldReturnAnyDocs());
+    invariant(PlanExecutor::NO_YIELD == request.getYieldPolicy());
+
+    MONGO_WRITE_CONFLICT_RETRY_LOOP_BEGIN {
+        // ParsedUpdate needs to be inside the write conflict retry loop because it may create a
+        // CanonicalQuery whose ownership will be transferred to the plan executor in
+        // getExecutorUpdate().
+        ParsedUpdate parsedUpdate(opCtx, &request);
+        auto parsedUpdateStatus = parsedUpdate.parseRequest();
+        if (!parsedUpdateStatus.isOK()) {
+            return parsedUpdateStatus;
+        }
+
+        AutoGetCollection autoColl(opCtx, nss, MODE_IX);
+        auto collectionResult = getCollection(
+            autoColl,
+            nss,
+            str::stream() << "Unable to update documents in " << nss.ns() << " using query "
+                          << query);
+        if (!collectionResult.isOK()) {
+            return collectionResult.getStatus();
+        }
+        auto collection = collectionResult.getValue();
+
+        auto planExecutorResult =
+            mongo::getExecutorUpdate(opCtx, nullptr, collection, &parsedUpdate);
+        if (!planExecutorResult.isOK()) {
+            return planExecutorResult.getStatus();
+        }
+        auto planExecutor = std::move(planExecutorResult.getValue());
+
+        return planExecutor->executePlan();
+    }
+    MONGO_WRITE_CONFLICT_RETRY_LOOP_END(opCtx, "_upsertWithQuery", nss.ns());
+
+    MONGO_UNREACHABLE;
+}
+
 }  // namespace
 
 Status StorageInterfaceImpl::upsertById(OperationContext* opCtx,
@@ -822,6 +894,12 @@ Status StorageInterfaceImpl::upsertById(OperationContext* opCtx,
     MONGO_UNREACHABLE;
 }
 
+Status StorageInterfaceImpl::putSingleton(OperationContext* opCtx,
+                                          const NamespaceString& nss,
+                                          const BSONObj& update) {
+    return _upsertWithQuery(opCtx, nss, {}, update);
+}
+
 Status StorageInterfaceImpl::deleteByFilter(OperationContext* opCtx,
                                             const NamespaceString& nss,
                                             const BSONObj& filter) {
diff --git a/src/mongo/db/repl/storage_interface_impl.h b/src/mongo/db/repl/storage_interface_impl.h
index f2f52ea4e5e..a41d0ae4584 100644
--- a/src/mongo/db/repl/storage_interface_impl.h
+++ b/src/mongo/db/repl/storage_interface_impl.h
@@ -123,6 +123,12 @@ public:
                                                      BoundInclusion boundInclusion,
                                                      std::size_t limit) override;
 
+    StatusWith<BSONObj> findSingleton(OperationContext* opCtx, const NamespaceString& nss) override;
+
+    Status putSingleton(OperationContext* opCtx,
+                        const NamespaceString& nss,
+                        const BSONObj& update) override;
+
     StatusWith<BSONObj> findById(OperationContext* opCtx,
                                  const NamespaceString& nss,
                                  const BSONElement& idKey) override;
diff --git a/src/mongo/db/repl/storage_interface_impl_test.cpp b/src/mongo/db/repl/storage_interface_impl_test.cpp
index ab06431d442..90cfd44f827 100644
--- a/src/mongo/db/repl/storage_interface_impl_test.cpp
+++ b/src/mongo/db/repl/storage_interface_impl_test.cpp
@@ -1619,6 +1619,109 @@ TEST_F(StorageInterfaceImplTest,
                       .getStatus());
 }
 
+TEST_F(StorageInterfaceImplTest, FindSingletonReturnsNamespaceNotFoundWhenDatabaseDoesNotExist) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    NamespaceString nss("nosuchdb.coll");
+    ASSERT_EQUALS(ErrorCodes::NamespaceNotFound, storage.findSingleton(opCtx, nss).getStatus());
+}
+
+TEST_F(StorageInterfaceImplTest, FindSingletonReturnsNamespaceNotFoundWhenCollectionDoesNotExist) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    NamespaceString nss("db.coll1");
+    ASSERT_OK(storage.createCollection(opCtx, NamespaceString("db.coll2"), CollectionOptions()));
+    ASSERT_EQUALS(ErrorCodes::NamespaceNotFound, storage.findSingleton(opCtx, nss).getStatus());
+}
+
+TEST_F(StorageInterfaceImplTest, FindSingletonReturnsCollectionIsEmptyWhenCollectionIsEmpty) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    auto nss = makeNamespace(_agent);
+    ASSERT_OK(storage.createCollection(opCtx, nss, CollectionOptions()));
+    ASSERT_EQUALS(ErrorCodes::CollectionIsEmpty, storage.findSingleton(opCtx, nss).getStatus());
+}
+
+TEST_F(StorageInterfaceImplTest,
+       FindSingletonReturnsTooManyMatchingDocumentsWhenNotSingletonCollection) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    auto nss = makeNamespace(_agent);
+    ASSERT_OK(storage.createCollection(opCtx, nss, CollectionOptions()));
+    auto doc1 = BSON("_id" << 0 << "x" << 0);
+    auto doc2 = BSON("_id" << 1 << "x" << 1);
+    ASSERT_OK(storage.insertDocuments(opCtx, nss, {doc1, doc2}));
+    ASSERT_EQUALS(ErrorCodes::TooManyMatchingDocuments,
+                  storage.findSingleton(opCtx, nss).getStatus());
+}
+
+TEST_F(StorageInterfaceImplTest, FindSingletonReturnsDocumentWhenSingletonDocumentExists) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    auto nss = makeNamespace(_agent);
+    ASSERT_OK(storage.createCollection(opCtx, nss, CollectionOptions()));
+    auto doc1 = BSON("_id" << 0 << "x" << 0);
+    ASSERT_OK(storage.insertDocument(opCtx, nss, doc1));
+    ASSERT_BSONOBJ_EQ(doc1, unittest::assertGet(storage.findSingleton(opCtx, nss)));
+}
+
+TEST_F(StorageInterfaceImplTest, PutSingletonReturnsNamespaceNotFoundWhenDatabaseDoesNotExist) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    NamespaceString nss("nosuchdb.coll");
+    auto update = BSON("$set" << BSON("_id" << 0 << "x" << 1));
+    ASSERT_EQUALS(ErrorCodes::NamespaceNotFound, storage.putSingleton(opCtx, nss, update));
+}
+
+TEST_F(StorageInterfaceImplTest, PutSingletonReturnsNamespaceNotFoundWhenCollectionDoesNotExist) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    NamespaceString nss("db.coll1");
+    ASSERT_OK(storage.createCollection(opCtx, nss, CollectionOptions()));
+    auto update = BSON("$set" << BSON("_id" << 0 << "x" << 1));
+    ASSERT_EQUALS(ErrorCodes::NamespaceNotFound,
+                  storage.putSingleton(opCtx, NamespaceString("db.coll2"), update));
+}
+
+TEST_F(StorageInterfaceImplTest, PutSingletonUpsertsDocumentsWhenCollectionIsEmpty) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    auto nss = makeNamespace(_agent);
+    ASSERT_OK(storage.createCollection(opCtx, nss, CollectionOptions()));
+    auto update = BSON("$set" << BSON("_id" << 0 << "x" << 1));
+    ASSERT_OK(storage.putSingleton(opCtx, nss, update));
+    ASSERT_BSONOBJ_EQ(BSON("_id" << 0 << "x" << 1),
+                      unittest::assertGet(storage.findSingleton(opCtx, nss)));
+    _assertDocumentsInCollectionEquals(opCtx, nss, {BSON("_id" << 0 << "x" << 1)});
+}
+
+TEST_F(StorageInterfaceImplTest, PutSingletonUpdatesDocumentWhenCollectionIsNotEmpty) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    auto nss = makeNamespace(_agent);
+    ASSERT_OK(storage.createCollection(opCtx, nss, CollectionOptions()));
+    auto doc1 = BSON("_id" << 0 << "x" << 0);
+    ASSERT_OK(storage.insertDocument(opCtx, nss, doc1));
+    auto update = BSON("$set" << BSON("x" << 1));
+    ASSERT_OK(storage.putSingleton(opCtx, nss, update));
+    ASSERT_BSONOBJ_EQ(BSON("_id" << 0 << "x" << 1),
+                      unittest::assertGet(storage.findSingleton(opCtx, nss)));
+    _assertDocumentsInCollectionEquals(opCtx, nss, {BSON("_id" << 0 << "x" << 1)});
+}
+
+TEST_F(StorageInterfaceImplTest, PutSingletonUpdatesFirstDocumentWhenCollectionIsNotSingleton) {
+    auto opCtx = getOperationContext();
+    StorageInterfaceImpl storage;
+    auto nss = makeNamespace(_agent);
+    ASSERT_OK(storage.createCollection(opCtx, nss, CollectionOptions()));
+    auto doc1 = BSON("_id" << 0 << "x" << 0);
+    auto doc2 = BSON("_id" << 1 << "x" << 1);
+    ASSERT_OK(storage.insertDocuments(opCtx, nss, {doc1, doc2}));
+    auto update = BSON("$set" << BSON("x" << 2));
+    ASSERT_OK(storage.putSingleton(opCtx, nss, update));
+    _assertDocumentsInCollectionEquals(opCtx, nss, {BSON("_id" << 0 << "x" << 2), doc2});
+}
+
 TEST_F(StorageInterfaceImplTest, FindByIdReturnsNamespaceNotFoundWhenDatabaseDoesNotExist) {
     auto opCtx = getOperationContext();
     StorageInterfaceImpl storage;
diff --git a/src/mongo/db/repl/storage_interface_mock.h b/src/mongo/db/repl/storage_interface_mock.h
index 17a17b20b87..eaf888b036e 100644
--- a/src/mongo/db/repl/storage_interface_mock.h
+++ b/src/mongo/db/repl/storage_interface_mock.h
@@ -202,6 +202,17 @@ public:
             opCtx, nss, indexName, scanDirection, startKey, boundInclusion, limit);
     }
 
+    StatusWith<BSONObj> findSingleton(OperationContext* opCtx,
+                                      const NamespaceString& nss) override {
+        return Status{ErrorCodes::IllegalOperation, "findSingleton not implemented."};
+    }
+
+    Status putSingleton(OperationContext* opCtx,
+                        const NamespaceString& nss,
+                        const BSONObj& update) override {
+        return Status{ErrorCodes::IllegalOperation, "putSingleton not implemented."};
+    }
+
     StatusWith<BSONObj> findById(OperationContext* opCtx,
                                  const NamespaceString& nss,
                                  const BSONElement& idKey) override {