Merge pull request #27062 from diogocarleto/mongoDBPanacheBugFixToAll…

…owCountWithCollation Support count with collation in MongoDB with Panache
quarkusio · Aug 2, 2022 · c3d0a84 · c3d0a84
2 parents f224e25 + d0bb072
commit c3d0a84
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 2 deletions.
diff --git a/...va/io/quarkus/mongodb/panache/common/reactive/runtime/CommonReactivePanacheQueryImpl.java b/...va/io/quarkus/mongodb/panache/common/reactive/runtime/CommonReactivePanacheQueryImpl.java
@@ -9,6 +9,7 @@
 
 import com.mongodb.ReadPreference;
 import com.mongodb.client.model.Collation;
+import com.mongodb.client.model.CountOptions;
 
 import io.quarkus.mongodb.FindOptions;
 import io.quarkus.mongodb.panache.common.runtime.MongoPropertyUtil;
@@ -154,7 +155,14 @@ public <T extends Entity> CommonReactivePanacheQueryImpl<T> withReadPreference(R
     @SuppressWarnings("unchecked")
     public Uni<Long> count() {
         if (count == null) {
-            count = mongoQuery == null ? collection.countDocuments() : collection.countDocuments(mongoQuery);
+            CountOptions countOptions = new CountOptions();
+            if (collation != null) {
+                countOptions.collation(collation);
+            }
+
+            count = mongoQuery == null
+                    ? collection.countDocuments()
+                    : collection.countDocuments(mongoQuery, countOptions);
         }
         return count;
     }

diff --git a/...ntime/src/main/java/io/quarkus/mongodb/panache/common/runtime/CommonPanacheQueryImpl.java b/...ntime/src/main/java/io/quarkus/mongodb/panache/common/runtime/CommonPanacheQueryImpl.java
@@ -16,6 +16,7 @@
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoCursor;
 import com.mongodb.client.model.Collation;
+import com.mongodb.client.model.CountOptions;
 
 import io.quarkus.panache.common.Page;
 import io.quarkus.panache.common.Range;
@@ -157,7 +158,14 @@ public <T extends Entity> CommonPanacheQueryImpl<T> withReadPreference(ReadPrefe
     public long count() {
         if (count == null) {
             Bson query = getQuery();
-            count = clientSession == null ? collection.countDocuments(query) : collection.countDocuments(clientSession, query);
+            CountOptions countOptions = new CountOptions();
+            if (collation != null) {
+                countOptions.collation(collation);
+            }
+
+            count = clientSession == null
+                    ? collection.countDocuments(query, countOptions)
+                    : collection.countDocuments(clientSession, query, countOptions);
         }
         return count;
     }

diff --git a/...-tests/mongodb-panache/src/main/java/io/quarkus/it/mongodb/panache/test/TestResource.java b/...-tests/mongodb-panache/src/main/java/io/quarkus/it/mongodb/panache/test/TestResource.java
@@ -113,6 +113,16 @@ public Response testImperativeEntity() {
         Assertions.assertEquals("aaa", results.get(0).title);
         Assertions.assertEquals("AAA", results.get(1).title);
         Assertions.assertEquals("BBB", results.get(2).title);
+
+        //count with collation
+        collation = Collation.builder()
+                .locale("en")
+                .collationStrength(CollationStrength.SECONDARY)
+                .build();
+        Assertions.assertEquals(2, TestImperativeEntity.find("{'title' : ?1}", "aaa").withCollation(collation).count());
+        Assertions.assertEquals(2, TestImperativeEntity.find("{'title' : ?1}", "AAA").withCollation(collation).count());
+        Assertions.assertEquals(1, TestImperativeEntity.find("{'title' : ?1}", "bbb").withCollation(collation).count());
+        Assertions.assertEquals(1, TestImperativeEntity.find("{'title' : ?1}", "BBB").withCollation(collation).count());
         entityAUpper.delete();
         entityALower.delete();
         entityB.delete();
@@ -264,6 +274,17 @@ public Response testImperativeRepository() {
         Assertions.assertEquals("aaa", results.get(0).title);
         Assertions.assertEquals("AAA", results.get(1).title);
         Assertions.assertEquals("BBB", results.get(2).title);
+
+        //count with collation
+        collation = Collation.builder()
+                .locale("en")
+                .collationStrength(CollationStrength.SECONDARY)
+                .build();
+
+        Assertions.assertEquals(2, testImperativeRepository.find("{'title' : ?1}", "aaa").withCollation(collation).count());
+        Assertions.assertEquals(2, testImperativeRepository.find("{'title' : ?1}", "AAA").withCollation(collation).count());
+        Assertions.assertEquals(1, testImperativeRepository.find("{'title' : ?1}", "bbb").withCollation(collation).count());
+        Assertions.assertEquals(1, testImperativeRepository.find("{'title' : ?1}", "BBB").withCollation(collation).count());
         testImperativeRepository.delete(entityALower);
         testImperativeRepository.delete(entityAUpper);
         testImperativeRepository.delete(entityB);
@@ -494,6 +515,20 @@ public Response testReactiveEntity() {
         Assertions.assertEquals("aaa", results.get(0).title);
         Assertions.assertEquals("AAA", results.get(1).title);
         Assertions.assertEquals("BBB", results.get(2).title);
+
+        //count with collation
+        collation = Collation.builder()
+                .locale("en")
+                .collationStrength(CollationStrength.SECONDARY)
+                .build();
+        Assertions.assertEquals(2, TestReactiveEntity.find("{'title' : ?1}", "aaa").withCollation(collation).count()
+                .await().indefinitely());
+        Assertions.assertEquals(2, TestReactiveEntity.find("{'title' : ?1}", "AAA").withCollation(collation).count()
+                .await().indefinitely());
+        Assertions.assertEquals(1, TestReactiveEntity.find("{'title' : ?1}", "bbb").withCollation(collation).count()
+                .await().indefinitely());
+        Assertions.assertEquals(1, TestReactiveEntity.find("{'title' : ?1}", "BBB").withCollation(collation).count()
+                .await().indefinitely());
         entityAUpper.delete().await().indefinitely();
         entityALower.delete().await().indefinitely();
         entityB.delete().await().indefinitely();
@@ -654,6 +689,20 @@ public Response testReactiveRepository() {
         Assertions.assertEquals("aaa", results.get(0).title);
         Assertions.assertEquals("AAA", results.get(1).title);
         Assertions.assertEquals("BBB", results.get(2).title);
+
+        //count with collation
+        collation = Collation.builder()
+                .locale("en")
+                .collationStrength(CollationStrength.SECONDARY)
+                .build();
+        Assertions.assertEquals(2, testReactiveRepository.find("{'title' : ?1}", "aaa").withCollation(collation).count()
+                .await().indefinitely());
+        Assertions.assertEquals(2, testReactiveRepository.find("{'title' : ?1}", "AAA").withCollation(collation).count()
+                .await().indefinitely());
+        Assertions.assertEquals(1, testReactiveRepository.find("{'title' : ?1}", "bbb").withCollation(collation).count()
+                .await().indefinitely());
+        Assertions.assertEquals(1, testReactiveRepository.find("{'title' : ?1}", "BBB").withCollation(collation).count()
+                .await().indefinitely());
         testReactiveRepository.delete(entityALower).await().indefinitely();
         testReactiveRepository.delete(entityAUpper).await().indefinitely();
         testReactiveRepository.delete(entityB).await().indefinitely();