[java-shell] support db.collection.find().count() (#450)

Author: kornilova203

packages/java-shell/src/main/kotlin/com/mongodb/mongosh/service/BaseMongoIterableHelper.kt

@@ -7,6 +7,7 @@ import com.mongodb.client.AggregateIterable
 import com.mongodb.client.FindIterable
 import com.mongodb.client.MongoDatabase
 import com.mongodb.client.MongoIterable
+import com.mongodb.client.model.CountOptions
 import com.mongodb.mongosh.MongoShellConverter
 import org.bson.Document
 import org.graalvm.polyglot.Value
@@ -39,7 +40,7 @@ internal abstract class BaseMongoIterableHelper<T : MongoIterable<*>>(val iterab
     open fun hint(v: Document): Unit = throw NotImplementedError("hint is not supported")
     open fun collation(v: Document): Unit = throw NotImplementedError("collation is not supported")
     open fun allowPartialResults(): Unit = throw NotImplementedError("allowPartialResults is not supported")
-    open fun count(): Int = throw NotImplementedError("count is not supported")
+    open fun count(): Long = throw NotImplementedError("count is not supported")
     open fun maxTimeMS(v: Long): Unit = throw NotImplementedError("maxTimeMS is not supported")
     open fun noCursorTimeout(): Unit = throw NotImplementedError("noCursorTimeout is not supported")
     open fun oplogReplay(): Unit = throw NotImplementedError("oplogReplay is not supported")
@@ -157,6 +158,14 @@ internal class FindIterableHelper(iterable: FindIterable<out Any?>,
     override fun sort(spec: Document) = set("sort", spec)
     override fun tailable() = set("tailable", CursorType.Tailable.toString())
 
+    override fun count(): Long {
+        check(createOptions != null) { "createOptions were not saved" }
+        val countOptionsMap = options.filterKeys { countOptionsConverters.containsKey(it) }
+        val countOptions = convert(CountOptions(), countOptionsConverters, countOptionsDefaultConverter, countOptionsMap).getOrThrow()
+        @Suppress("DEPRECATION")
+        return createOptions.db.getCollection(createOptions.collection).count(createOptions.find, countOptions)
+    }
+
     override fun explain(verbosity: String?): Any? {
         set("explain", true)
         return iterable.first()

packages/java-shell/src/main/kotlin/com/mongodb/mongosh/service/Cursor.kt

@@ -49,6 +49,9 @@ internal class Cursor(private var helper: BaseMongoIterableHelper<*>, private va
         return this
     }
 
+    /**
+     * cursor.objsLeftInBatch()
+     */
     override fun bufferedCount(): Int {
         throw NotImplementedError("bufferedCount is not supported")
     }
@@ -89,7 +92,7 @@ internal class Cursor(private var helper: BaseMongoIterableHelper<*>, private va
     }
 
     @HostAccess.Export
-    override fun count(): Int {
+    override fun count(): Long {
         checkQueryNotExecuted()
         return helper.count()
     }

packages/java-shell/src/main/kotlin/com/mongodb/mongosh/service/ServiceProviderCursor.kt

@@ -13,7 +13,7 @@ interface ServiceProviderCursor {
   fun isClosed(): Boolean
   fun collation(v: Value): ServiceProviderCursor
   fun comment(v: String): ServiceProviderCursor
-  fun count(): Int
+  fun count(): Long
   fun forEach(func: Value)
   fun hasNext(): Boolean
   fun hint(v: Value): ServiceProviderCursor

packages/java-shell/src/test/kotlin/com/mongodb/mongosh/CollectionTest.kt

@@ -32,6 +32,7 @@ class CollectionTest : ShellTestCase() {
     @Test fun testEstimatedDocumentCount() = test()
     @Test fun testEstimatedDocumentCountWithMaxTime() = test()
     @Test fun testFind()                   = test()
+    @Test fun testFindAndCount()           = test()
     @Test fun testFindMap()                = test()
     @Test fun testFindOne()                = test()
     @Test fun testFindOneAndDelete()       = test()

packages/java-shell/src/test/resources/collection/findAndCount.expected.txt

@@ -0,0 +1,3 @@
+1
+1
+2

packages/java-shell/src/test/resources/collection/findAndCount.js

@@ -0,0 +1,11 @@
+// before
+db.coll.deleteMany({});
+db.coll.insertMany([{v: 1}, {v: 2}]);
+// command
+db.coll.find({v: 1}).count();
+// command
+db.coll.find().limit(1).count();
+// command
+db.coll.find({}, {}, {"comment": "hello"}).count();
+// clear
+db.coll.drop();