From 475b55c3fb074ff1d7f97c685f1340cacc685f2c Mon Sep 17 00:00:00 2001
From: Karl F Walkow <karl@walkow.de>
Date: Mon, 24 Jun 2024 12:09:53 +0200
Subject: [PATCH] Support counting documents with options

---
 .../src/main/scala/molly/core/MollyCollection.scala |  6 ++++++
 .../src/main/scala/molly/core/syntax/model.scala    |  5 +++++
 .../test/scala/molly/core/MollyCollectionTest.scala | 13 +++++++++++++
 3 files changed, 24 insertions(+)

diff --git a/molly-core/src/main/scala/molly/core/MollyCollection.scala b/molly-core/src/main/scala/molly/core/MollyCollection.scala
index ca61613..ca1e3ee 100644
--- a/molly-core/src/main/scala/molly/core/MollyCollection.scala
+++ b/molly-core/src/main/scala/molly/core/MollyCollection.scala
@@ -21,6 +21,7 @@ import molly.core.reactivestreams.fromSinglePublisher
 import molly.core.reactivestreams.fromStreamPublisher
 import molly.core.reactivestreams.fromVoidPublisher
 import molly.core.syntax.model.BulkWriteOptions
+import molly.core.syntax.model.CountOptions
 import molly.core.syntax.model.CreateIndexOptions
 import molly.core.syntax.model.FindOneAndReplaceOptions
 import molly.core.syntax.model.FindOneAndUpdateOptions
@@ -64,6 +65,11 @@ final case class MollyCollection[F[_], A] private[core] (private[core] val deleg
     */
   def countDocuments(filter: Bson): F[Long] = fromSinglePublisher(delegate.countDocuments(filter)).map(_.toLong)
 
+  /** [[https://mongodb.github.io/mongo-java-driver/5.1/apidocs/mongodb-driver-reactivestreams/com/mongodb/reactivestreams/client/MongoCollection.html#countDocuments(org.bson.conversions.Bson,com.mongodb.client.model.CountOptions)]]
+    */
+  def countDocuments(filter: Bson, options: CountOptions): F[Long] =
+    fromSinglePublisher(delegate.countDocuments(filter, options)).map(_.toLong)
+
   /** [[https://mongodb.github.io/mongo-java-driver/5.1/apidocs/mongodb-driver-reactivestreams/com/mongodb/reactivestreams/client/MongoCollection.html#createIndex(org.bson.conversions.Bson)]]
     */
   def createIndex(key: Bson): F[String] = fromSinglePublisher(delegate.createIndex(key))
diff --git a/molly-core/src/main/scala/molly/core/syntax/model.scala b/molly-core/src/main/scala/molly/core/syntax/model.scala
index 03e7a8f..290d0f2 100644
--- a/molly-core/src/main/scala/molly/core/syntax/model.scala
+++ b/molly-core/src/main/scala/molly/core/syntax/model.scala
@@ -13,6 +13,11 @@ trait model:
   object BulkWriteOptions:
     def apply(): BulkWriteOptions = new BulkWriteOptions()
 
+  type CountOptions = com.mongodb.client.model.CountOptions
+
+  object CountOptions:
+    def apply(): CountOptions = new CountOptions()
+
   type CreateIndexOptions = com.mongodb.client.model.CreateIndexOptions
 
   object CreateIndexOptions:
diff --git a/molly-core/src/test/scala/molly/core/MollyCollectionTest.scala b/molly-core/src/test/scala/molly/core/MollyCollectionTest.scala
index 5424789..518f3ff 100644
--- a/molly-core/src/test/scala/molly/core/MollyCollectionTest.scala
+++ b/molly-core/src/test/scala/molly/core/MollyCollectionTest.scala
@@ -11,6 +11,7 @@ import com.mongodb.client.model.ReturnDocument
 import com.mongodb.client.model.Sorts
 import com.mongodb.client.model.Updates
 import molly.core.syntax.bsondocument.BsonDocumentCollection
+import molly.core.syntax.model.CountOptions
 import molly.core.syntax.model.CreateIndexOptions
 import molly.core.syntax.model.DeleteOneModel
 import molly.core.syntax.model.FindOneAndReplaceOptions
@@ -112,6 +113,18 @@ object MollyCollectionTest extends IOSuite with TestContainerForAll[IO] with Mol
         count <- coll.countDocuments(Filters.regex("foo", "bar.*"))
       yield expect(count == 2L)
 
+  test("countDocuments: count documents given a filter and options"): containers =>
+    withClient(containers): (client: MollyClient[IO]) =>
+      val doc1 = new BsonDocument("foo", new BsonString("bar"))
+      val doc2 = new BsonDocument("foo", new BsonString("baz"))
+      val doc3 = new BsonDocument("foo", new BsonString("barry"))
+      for
+        db    <- client.getDatabase("test")
+        coll  <- db.getCollection("countDocuments2")
+        _     <- coll.insertMany(Seq(doc1, doc2, doc3))
+        count <- coll.countDocuments(Filters.regex("foo", "bar.*"), CountOptions().limit(1))
+      yield expect(count == 1L)
+
   test("create and list index"): containers =>
     withClient(containers): (client: MollyClient[IO]) =>
       for