Always use count command for estimatedDocumentCount (mongodb#929)

JAVA-4530
neerajbhatt · May 4, 2022 · 4e0b916 · 4e0b916
1 parent 1ffcc02
commit 4e0b916
Show file tree

Hide file tree

Showing 15 changed files with 196 additions and 1,571 deletions.
diff --git a/driver-core/src/main/com/mongodb/internal/operation/EstimatedDocumentCountOperation.java b/driver-core/src/main/com/mongodb/internal/operation/EstimatedDocumentCountOperation.java
@@ -22,13 +22,10 @@
 import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.internal.binding.AsyncReadBinding;
 import com.mongodb.internal.binding.ReadBinding;
-import com.mongodb.internal.connection.QueryResult;
 import com.mongodb.internal.operation.CommandOperationHelper.CommandReadTransformer;
 import com.mongodb.internal.operation.CommandOperationHelper.CommandReadTransformerAsync;
 import com.mongodb.internal.session.SessionContext;
-import org.bson.BsonArray;
 import org.bson.BsonDocument;
-import org.bson.BsonInt32;
 import org.bson.BsonString;
 import org.bson.codecs.BsonDocumentCodec;
 import org.bson.codecs.Decoder;
@@ -42,11 +39,8 @@
 import static com.mongodb.internal.operation.CommandOperationHelper.isNamespaceError;
 import static com.mongodb.internal.operation.CommandOperationHelper.rethrowIfNotNamespaceError;
 import static com.mongodb.internal.operation.DocumentHelper.putIfNotZero;
-import static com.mongodb.internal.operation.OperationHelper.cursorDocumentToQueryResult;
 import static com.mongodb.internal.operation.OperationHelper.validateReadConcern;
 import static com.mongodb.internal.operation.OperationReadConcernHelper.appendReadConcernToCommand;
-import static com.mongodb.internal.operation.ServerVersionHelper.serverIsAtLeastVersionFiveDotZero;
-import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
 
 public class EstimatedDocumentCountOperation implements AsyncReadOperation<Long>, ReadOperation<Long> {
@@ -102,45 +96,16 @@ private CommandReadTransformerAsync<BsonDocument, Long> asyncTransformer() {
     }
 
     private long transformResult(final BsonDocument result, final ConnectionDescription connectionDescription) {
-        if (serverIsAtLeastVersionFiveDotZero(connectionDescription)) {
-            QueryResult<BsonDocument> queryResult = cursorDocumentToQueryResult(result.getDocument("cursor"),
-                    connectionDescription.getServerAddress());
-            return queryResult.getResults().get(0).getNumber("n").longValue();
-        } else {
-            return (result.getNumber("n")).longValue();
-        }
+        return (result.getNumber("n")).longValue();
     }
 
     private CommandCreator getCommandCreator(final SessionContext sessionContext) {
         return (serverDescription, connectionDescription) -> {
-            if (serverIsAtLeastVersionFiveDotZero(connectionDescription)) {
-                return getAggregateCommand(sessionContext, connectionDescription);
-            } else {
-                validateReadConcern(connectionDescription, sessionContext.getReadConcern());
-                return getCountCommand(sessionContext, connectionDescription);
-            }
+            validateReadConcern(connectionDescription, sessionContext.getReadConcern());
+            BsonDocument document = new BsonDocument("count", new BsonString(namespace.getCollectionName()));
+            appendReadConcernToCommand(sessionContext, connectionDescription.getMaxWireVersion(), document);
+            putIfNotZero(document, "maxTimeMS", maxTimeMS);
+            return document;
         };
     }
-
-    private BsonDocument getAggregateCommand(final SessionContext sessionContext, final ConnectionDescription connectionDescription) {
-        BsonDocument document = new BsonDocument("aggregate", new BsonString(namespace.getCollectionName()))
-                .append("cursor", new BsonDocument())
-                .append("pipeline", new BsonArray(asList(
-                     new BsonDocument("$collStats", new BsonDocument("count", new BsonDocument())),
-                     new BsonDocument("$group", new BsonDocument("_id", new BsonInt32(1))
-                             .append("n", new BsonDocument("$sum", new BsonString("$count")))
-                ))));
-
-        appendReadConcernToCommand(sessionContext, connectionDescription.getMaxWireVersion(), document);
-        putIfNotZero(document, "maxTimeMS", maxTimeMS);
-        return document;
-    }
-
-    private BsonDocument getCountCommand(final SessionContext sessionContext, final ConnectionDescription connectionDescription) {
-        BsonDocument document = new BsonDocument("count", new BsonString(namespace.getCollectionName()));
-
-        appendReadConcernToCommand(sessionContext, connectionDescription.getMaxWireVersion(), document);
-        putIfNotZero(document, "maxTimeMS", maxTimeMS);
-        return document;
-    }
 }
diff --git a/driver-core/src/test/resources/atlas-data-lake/estimatedDocumentCount.json b/driver-core/src/test/resources/atlas-data-lake/estimatedDocumentCount.json
@@ -15,24 +15,9 @@
         {
           "command_started_event": {
             "command": {
-              "aggregate": "driverdata",
-              "pipeline": [
-                {
-                  "$collStats": {
-                    "count": {}
-                  }
-                },
-                {
-                  "$group": {
-                    "_id": 1,
-                    "n": {
-                      "$sum": "$count"
-                    }
-                  }
-                }
-              ]
+              "count": "driverdata"
             },
-            "command_name": "aggregate",
+            "command_name": "count",
             "database_name": "test"
           }
         }

diff --git a/driver-core/src/test/resources/retryable-reads/estimatedDocumentCount-4.9.json b/driver-core/src/test/resources/retryable-reads/estimatedDocumentCount-4.9.json