feat: mongdb read preferences by query

quarkusio · Dec 10, 2020 · a0954ec · a0954ec
1 parent de1b9c6
commit a0954ec
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 2 deletions.
diff --git a/...mmon/runtime/src/main/java/io/quarkus/mongodb/panache/runtime/CommonPanacheQueryImpl.java b/...mmon/runtime/src/main/java/io/quarkus/mongodb/panache/runtime/CommonPanacheQueryImpl.java
@@ -9,6 +9,7 @@
 import org.bson.Document;
 import org.bson.conversions.Bson;
 
+import com.mongodb.ReadPreference;
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoCursor;
@@ -30,6 +31,7 @@ public class CommonPanacheQueryImpl<Entity> {
     private Range range;
 
     private Collation collation;
+    private ReadPreference readPreference;
 
     public CommonPanacheQueryImpl(MongoCollection<? extends Entity> collection, Bson mongoQuery, Bson sort) {
         this.collection = collection;
@@ -46,6 +48,7 @@ private CommonPanacheQueryImpl(CommonPanacheQueryImpl<?> previousQuery, Bson pro
         this.count = previousQuery.count;
         this.range = previousQuery.range;
         this.collation = previousQuery.collation;
+        this.readPreference = previousQuery.readPreference;
     }
 
     public <T> CommonPanacheQueryImpl<T> project(Class<T> type) {
@@ -139,12 +142,21 @@ public <T extends Entity> CommonPanacheQueryImpl<T> withCollation(Collation coll
         return (CommonPanacheQueryImpl<T>) this;
     }
 
+    public <T extends Entity> CommonPanacheQueryImpl<T> withReadPreference(ReadPreference readPreference) {
+        this.readPreference = readPreference;
+        return (CommonPanacheQueryImpl<T>) this;
+    }
+
     // Results
 
     @SuppressWarnings("unchecked")
     public long count() {
         if (count == null) {
-            count = collection.countDocuments(mongoQuery);
+            MongoCollection theCollection = collection;
+            if (this.readPreference != null) {
+                theCollection = theCollection.withReadPreference(this.readPreference);
+            }
+            count = theCollection.countDocuments(mongoQuery);
         }
         return count;
     }
@@ -156,7 +168,11 @@ public <T extends Entity> List<T> list() {
     @SuppressWarnings("unchecked")
     private <T extends Entity> List<T> list(Integer limit) {
         List<T> list = new ArrayList<>();
-        FindIterable find = mongoQuery == null ? collection.find() : collection.find(mongoQuery);
+        MongoCollection theCollection = collection;
+        if (this.readPreference != null) {
+            theCollection = theCollection.withReadPreference(this.readPreference);
+        }
+        FindIterable find = mongoQuery == null ? theCollection.find() : theCollection.find(mongoQuery);
         if (this.projections != null) {
             find.projection(projections);
         }

diff --git a/...lin/runtime/src/main/java/io/quarkus/mongodb/panache/kotlin/runtime/PanacheQueryImpl.java b/...lin/runtime/src/main/java/io/quarkus/mongodb/panache/kotlin/runtime/PanacheQueryImpl.java
@@ -5,6 +5,7 @@
 
 import org.bson.conversions.Bson;
 
+import com.mongodb.ReadPreference;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.model.Collation;
 
@@ -95,6 +96,12 @@ public PanacheQuery<Entity> withCollation(Collation collation) {
         return this;
     }
 
+    @Override
+    public PanacheQuery<Entity> withReadPreference(ReadPreference readPreference) {
+        delegate.withReadPreference(readPreference);
+        return this;
+    }
+
     // Results
 
     @Override

diff --git a/...-panache-kotlin/runtime/src/main/kotlin/io/quarkus/mongodb/panache/kotlin/PanacheQuery.kt b/...-panache-kotlin/runtime/src/main/kotlin/io/quarkus/mongodb/panache/kotlin/PanacheQuery.kt
@@ -1,5 +1,6 @@
 package io.quarkus.mongodb.panache.kotlin
 
+import com.mongodb.ReadPreference
 import com.mongodb.client.model.Collation
 import io.quarkus.panache.common.Page
 import io.quarkus.panache.common.exception.PanacheQueryException
@@ -137,6 +138,15 @@ interface PanacheQuery<Entity: Any> {
      * @return this query, modified
      */
     fun withCollation(collation: Collation): PanacheQuery<Entity>
+
+    /**
+     * Define the read preference used for this query.
+     *
+     * @param readPreference the read preference to be used for this query.
+     * @return this query, modified
+     */
+    fun withReadPreference(readPreference: ReadPreference?): PanacheQuery<Entity>
+
     // Results
     /**
      * Reads and caches the total number of entities this query operates on. This causes a database

diff --git a/...anache/mongodb-panache/runtime/src/main/java/io/quarkus/mongodb/panache/PanacheQuery.java b/...anache/mongodb-panache/runtime/src/main/java/io/quarkus/mongodb/panache/PanacheQuery.java
@@ -4,6 +4,7 @@
 import java.util.Optional;
 import java.util.stream.Stream;
 
+import com.mongodb.ReadPreference;
 import com.mongodb.client.model.Collation;
 
 import io.quarkus.panache.common.Page;
@@ -144,6 +145,14 @@ public interface PanacheQuery<Entity> {
      */
     public <T extends Entity> PanacheQuery<T> withCollation(Collation collation);
 
+    /**
+     * Define the read preference used for this query.
+     *
+     * @param readPreference the read preference to be used for this query.
+     * @return this query, modified
+     */
+    public <T extends Entity> PanacheQuery<T> withReadPreference(ReadPreference readPreference);
+
     // Results
 
     /**

diff --git a/...db-panache/runtime/src/main/java/io/quarkus/mongodb/panache/runtime/PanacheQueryImpl.java b/...db-panache/runtime/src/main/java/io/quarkus/mongodb/panache/runtime/PanacheQueryImpl.java
@@ -6,6 +6,7 @@
 
 import org.bson.conversions.Bson;
 
+import com.mongodb.ReadPreference;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.model.Collation;
 
@@ -97,6 +98,12 @@ public <T extends Entity> PanacheQuery<T> withCollation(Collation collation) {
         return (PanacheQuery<T>) this;
     }
 
+    @Override
+    public <T extends Entity> PanacheQuery<T> withReadPreference(ReadPreference readPreference) {
+        delegate.withReadPreference(readPreference);
+        return (PanacheQuery<T>) this;
+    }
+
     // Results
 
     @Override