SERVER-36094 Make query by UUID work for DBClient.

7 files changed, 144 insertions, 6 deletions

diff --git a/src/mongo/client/dbclient_base.h b/src/mongo/client/dbclient_base.h
index 88c5cd1de83..13c8de941ef 100644
--- a/src/mongo/client/dbclient_base.h
+++ b/src/mongo/client/dbclient_base.h
@@ -595,6 +595,10 @@ public:
 
         Use the DBClientCursorBatchIterator version, below, if you want to do items in large
         blocks, perhaps to avoid granular locking and such.
+
+        Note:
+        The version that takes a BSONObj cannot return the namespace queried when the query is
+        is done by UUID.  If this is required, use the DBClientBatchIterator version.
      */
     unsigned long long query(stdx::function<void(const BSONObj&)> f,
                              const NamespaceStringOrUUID& nsOrUuid,
diff --git a/src/mongo/client/dbclient_cursor.cpp b/src/mongo/client/dbclient_cursor.cpp
index ff5698acc00..b16184c9700 100644
--- a/src/mongo/client/dbclient_cursor.cpp
+++ b/src/mongo/client/dbclient_cursor.cpp
@@ -111,14 +111,19 @@ Message DBClientCursor::_assembleInit() {
             return assembleCommandRequest(_client, ns.db(), opts, query);
         }
     } else if (_useFindCommand) {
-        auto qr = QueryRequest::fromLegacyQuery(ns,
+        // The caller supplies a 'query' object which may have $-prefixed directives in the format
+        // expected for a legacy OP_QUERY. Therefore, we use the legacy parsing code supplied by
+        // QueryRequest. When actually issuing the request to the remote node, we will assemble a
+        // find command.
+        auto qr = QueryRequest::fromLegacyQuery(_nsOrUuid,
                                                 query,
                                                 fieldsToReturn ? *fieldsToReturn : BSONObj(),
                                                 nToSkip,
                                                 nextBatchSize(),
                                                 opts);
         if (qr.isOK() && !qr.getValue()->isExplain()) {
-            BSONObj cmd = qr.getValue()->asFindCommand();
+            BSONObj cmd = _nsOrUuid.uuid() ? qr.getValue()->asFindCommandWithUuid()
+                                           : qr.getValue()->asFindCommand();
             if (auto readPref = query["$readPreference"]) {
                 // QueryRequest doesn't handle $readPreference.
                 cmd = BSONObjBuilder(std::move(cmd)).append(readPref).obj();
@@ -126,6 +131,13 @@ Message DBClientCursor::_assembleInit() {
             return assembleCommandRequest(_client, ns.db(), opts, std::move(cmd));
         }
         // else use legacy OP_QUERY request.
+        // Legacy OP_QUERY request does not support UUIDs.
+        if (_nsOrUuid.uuid()) {
+            // If there was a problem building the query request, report that.
+            uassertStatusOK(qr.getStatus());
+            // Otherwise it must have been explain.
+            uasserted(50937, "Query by UUID is not supported for explain queries.");
+        }
     }
 
     _useFindCommand = false;  // Make sure we handle the reply correctly.
@@ -522,6 +534,7 @@ DBClientCursor::DBClientCursor(DBClientBase* client,
     : batch{std::move(initialBatch)},
       _client(client),
       _originalHost(_client->getServerAddress()),
+      _nsOrUuid(nsOrUuid),
       ns(nsOrUuid.nss() ? *nsOrUuid.nss() : NamespaceString(nsOrUuid.dbname())),
       _isCommand(ns.isCommand()),
       query(query),
@@ -536,8 +549,11 @@ DBClientCursor::DBClientCursor(DBClientBase* client,
       _ownCursor(true),
       wasError(false),
       _enabledBSONVersion(Validator<BSONObj>::enabledBSONVersion()) {
-    if (queryOptions & QueryOptionLocal_forceOpQuery)
+    if (queryOptions & QueryOptionLocal_forceOpQuery) {
+        // Legacy OP_QUERY does not support UUIDs.
+        invariant(!_nsOrUuid.uuid());
         _useFindCommand = false;
+    }
 }
 
 DBClientCursor::~DBClientCursor() {
diff --git a/src/mongo/client/dbclient_cursor.h b/src/mongo/client/dbclient_cursor.h
index 7b0f017238b..bf0c520552c 100644
--- a/src/mongo/client/dbclient_cursor.h
+++ b/src/mongo/client/dbclient_cursor.h
@@ -180,6 +180,9 @@ public:
         return ns.ns();
     }
 
+    const NamespaceString& getNamespaceString() const {
+        return ns;
+    }
     /**
      * actually does the query
      */
@@ -244,6 +247,10 @@ private:
     Batch batch;
     DBClientBase* _client;
     std::string _originalHost;
+    NamespaceStringOrUUID _nsOrUuid;
+    // 'ns' is initially the NamespaceString passed in, or the dbName if doing a find by UUID.
+    // After a successful 'find' command, 'ns' is updated to contain the namespace returned by that
+    // command.
     NamespaceString ns;
     const bool _isCommand;
     BSONObj query;
@@ -301,6 +308,10 @@ public:
     int n() const {
         return _n;
     }
+    // getNamespaceString() will return the NamespaceString returned by the 'find' command.
+    const NamespaceString& getNamespaceString() {
+        return _c.getNamespaceString();
+    }
 
 private:
     DBClientCursor& _c;
diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp
index 55dc81482e5..86df62f340d 100644
--- a/src/mongo/db/query/query_request.cpp
+++ b/src/mongo/db/query/query_request.cpp
@@ -408,9 +408,24 @@ BSONObj QueryRequest::asFindCommand() const {
     return bob.obj();
 }
 
+BSONObj QueryRequest::asFindCommandWithUuid() const {
+    BSONObjBuilder bob;
+    asFindCommandWithUuid(&bob);
+    return bob.obj();
+}
+
 void QueryRequest::asFindCommand(BSONObjBuilder* cmdBuilder) const {
     cmdBuilder->append(kFindCommandName, _nss.coll());
+    asFindCommandInternal(cmdBuilder);
+}
+
+void QueryRequest::asFindCommandWithUuid(BSONObjBuilder* cmdBuilder) const {
+    invariant(_uuid);
+    _uuid->appendToBuilder(cmdBuilder, kFindCommandName);
+    asFindCommandInternal(cmdBuilder);
+}
 
+void QueryRequest::asFindCommandInternal(BSONObjBuilder* cmdBuilder) const {
     if (!_filter.isEmpty()) {
         cmdBuilder->append(kFilterField, _filter);
     }
@@ -710,13 +725,13 @@ StatusWith<unique_ptr<QueryRequest>> QueryRequest::fromLegacyQueryMessage(const
     return std::move(qr);
 }
 
-StatusWith<unique_ptr<QueryRequest>> QueryRequest::fromLegacyQuery(NamespaceString nss,
+StatusWith<unique_ptr<QueryRequest>> QueryRequest::fromLegacyQuery(NamespaceStringOrUUID nsOrUuid,
                                                                    const BSONObj& queryObj,
                                                                    const BSONObj& proj,
                                                                    int ntoskip,
                                                                    int ntoreturn,
                                                                    int queryOptions) {
-    auto qr = stdx::make_unique<QueryRequest>(nss);
+    auto qr = stdx::make_unique<QueryRequest>(nsOrUuid);
 
     Status status = qr->init(ntoskip, ntoreturn, queryOptions, queryObj, proj, true);
     if (!status.isOK()) {
diff --git a/src/mongo/db/query/query_request.h b/src/mongo/db/query/query_request.h
index ccf612b89ae..0e56a033d4a 100644
--- a/src/mongo/db/query/query_request.h
+++ b/src/mongo/db/query/query_request.h
@@ -83,9 +83,12 @@ public:
 
     /**
      * Converts this QR into a find command.
+     * The withUuid variants make a UUID-based find command instead of a namespace-based ones.
      */
     BSONObj asFindCommand() const;
+    BSONObj asFindCommandWithUuid() const;
     void asFindCommand(BSONObjBuilder* cmdBuilder) const;
+    void asFindCommandWithUuid(BSONObjBuilder* cmdBuilder) const;
 
     /**
      * Converts this QR into an aggregation using $match. If this QR has options that cannot be
@@ -391,7 +394,7 @@ public:
     /**
      * Parse the provided legacy query object and parameters to construct a QueryRequest.
      */
-    static StatusWith<std::unique_ptr<QueryRequest>> fromLegacyQuery(NamespaceString nss,
+    static StatusWith<std::unique_ptr<QueryRequest>> fromLegacyQuery(NamespaceStringOrUUID nsOrUuid,
                                                                      const BSONObj& queryObj,
                                                                      const BSONObj& proj,
                                                                      int ntoskip,
@@ -432,6 +435,11 @@ private:
      */
     void addMetaProjection();
 
+    /**
+     * Common code for UUID and namespace-based find commands.
+     */
+    void asFindCommandInternal(BSONObjBuilder* cmdBuilder) const;
+
     NamespaceString _nss;
     OptionalCollectionUUID _uuid;
 
diff --git a/src/mongo/db/query/query_request_test.cpp b/src/mongo/db/query/query_request_test.cpp
index 1e3822b10ba..b83e89f0b6c 100644
--- a/src/mongo/db/query/query_request_test.cpp
+++ b/src/mongo/db/query/query_request_test.cpp
@@ -851,6 +851,60 @@ TEST(QueryRequestTest, ParseFromCommandDefaultBatchSize) {
 }
 
 //
+// Test asFindCommand ns and uuid variants.
+//
+
+TEST(QueryRequestTest, AsFindCommandAllNonOptionFields) {
+    BSONObj cmdObj = fromjson(
+        "{find: 'testns',"
+        "filter: {a: 1},"
+        "projection: {c: 1},"
+        "sort: {b: 1},"
+        "hint: {d: 1},"
+        "readConcern: {e: 1},"
+        "collation: {f: 1},"
+        "skip: 5,"
+        "limit: 3,"
+        "batchSize: 90,"
+        "singleBatch: true}");
+    const NamespaceString nss("test.testns");
+    bool isExplain = false;
+    unique_ptr<QueryRequest> qr(
+        assertGet(QueryRequest::makeFromFindCommand(nss, cmdObj, isExplain)));
+    ASSERT_BSONOBJ_EQ(cmdObj, qr->asFindCommand());
+}
+
+TEST(QueryRequestTest, AsFindCommandWithUuidAllNonOptionFields) {
+    BSONObj cmdObj = fromjson(
+        // This binary value is UUID("01234567-89ab-cdef-edcb-a98765432101")
+        "{find: { \"$binary\" : \"ASNFZ4mrze/ty6mHZUMhAQ==\", \"$type\" : \"04\" },"
+        "filter: {a: 1},"
+        "projection: {c: 1},"
+        "sort: {b: 1},"
+        "hint: {d: 1},"
+        "readConcern: {e: 1},"
+        "collation: {f: 1},"
+        "skip: 5,"
+        "limit: 3,"
+        "batchSize: 90,"
+        "singleBatch: true}");
+    const NamespaceString nss("test.testns");
+    bool isExplain = false;
+    unique_ptr<QueryRequest> qr(
+        assertGet(QueryRequest::makeFromFindCommand(nss, cmdObj, isExplain)));
+    ASSERT_BSONOBJ_EQ(cmdObj, qr->asFindCommandWithUuid());
+}
+
+TEST(QueryRequestTest, AsFindCommandWithUuidNoAvailableNamespace) {
+    BSONObj cmdObj =
+        fromjson("{find: { \"$binary\" : \"ASNFZ4mrze/ty6mHZUMhAQ==\", \"$type\" : \"04\" }}");
+    QueryRequest qr(NamespaceStringOrUUID(
+        "test", UUID::parse("01234567-89ab-cdef-edcb-a98765432101").getValue()));
+    ASSERT_BSONOBJ_EQ(cmdObj, qr.asFindCommandWithUuid());
+}
+
+//
+//
 // Errors checked in QueryRequest::validate().
 //
 
diff --git a/src/mongo/dbtests/querytests.cpp b/src/mongo/dbtests/querytests.cpp
index 239ceae45ab..1193a26dd08 100644
--- a/src/mongo/dbtests/querytests.cpp
+++ b/src/mongo/dbtests/querytests.cpp
@@ -1668,6 +1668,35 @@ public:
     }
 };
 
+class QueryByUuid : public CollectionBase {
+public:
+    QueryByUuid() : CollectionBase("QueryByUuid") {}
+
+    void run() {
+        CollectionOptions coll_opts;
+        coll_opts.uuid = UUID::gen();
+        {
+            Lock::GlobalWrite lk(&_opCtx);
+            OldClientContext context(&_opCtx, ns());
+            WriteUnitOfWork wunit(&_opCtx);
+            context.db()->createCollection(&_opCtx, ns(), coll_opts, false);
+            wunit.commit();
+        }
+        insert(ns(), BSON("a" << 1));
+        insert(ns(), BSON("a" << 2));
+        insert(ns(), BSON("a" << 3));
+        unique_ptr<DBClientCursor> cursor =
+            _client.query(NamespaceStringOrUUID("unittests", *coll_opts.uuid), BSONObj());
+        ASSERT_EQUALS(string(ns()), cursor->getns());
+        for (int i = 1; i <= 3; ++i) {
+            ASSERT(cursor->more());
+            BSONObj obj(cursor->next());
+            ASSERT_EQUALS(obj["a"].Int(), i);
+        }
+        ASSERT(!cursor->more());
+    }
+};
+
 class CollectionInternalBase : public CollectionBase {
 public:
     CollectionInternalBase(const char* nsLeaf)
@@ -1833,6 +1862,7 @@ public:
         add<FindingStartPartiallyFull>();
         add<FindingStartStale>();
         add<WhatsMyUri>();
+        add<QueryByUuid>();
         add<Exhaust>();
         add<QueryReadsAll>();
         add<queryobjecttests::names1>();