Simplify options timeout setting

mongodb · Apr 5, 2024 · aaa1ebd · aaa1ebd
1 parent 3807b89
commit aaa1ebd
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/driver-sync/src/main/com/mongodb/client/internal/ClientEncryptionImpl.java b/driver-sync/src/main/com/mongodb/client/internal/ClientEncryptionImpl.java
@@ -66,6 +66,7 @@
 import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.MINUTES;
 import static org.bson.internal.BsonUtil.mutableDeepCopy;
 
 /**
@@ -143,17 +144,17 @@ public BsonValue decrypt(final BsonBinary value) {
 
     @Override
     public DeleteResult deleteKey(final BsonBinary id) {
-        return collectionWithTimeout(collection, startTimeout()).deleteOne(Filters.eq("_id", id));
+        return collectionWithOptionsTimeout().deleteOne(Filters.eq("_id", id));
     }
 
     @Override
     public BsonDocument getKey(final BsonBinary id) {
-        return collectionWithTimeout(collection, startTimeout()).find(Filters.eq("_id", id)).first();
+        return collectionWithOptionsTimeout().find(Filters.eq("_id", id)).first();
     }
 
     @Override
     public FindIterable<BsonDocument> getKeys() {
-        return collectionWithTimeout(collection, startTimeout()).find();
+        return collectionWithTimeout().find();
     }
 
     @Override
@@ -283,4 +284,9 @@ public void close() {
     private Timeout startTimeout() {
         return TimeoutContext.calculateTimeout(options.getTimeout(MILLISECONDS));
     }
+
+    private MongoCollection<BsonDocument> collectionWithOptionsTimeout() {
+        Long timeout = options.getTimeout(MILLISECONDS);
+        return timeout != null ? collection.withTimeout(timeout, MILLISECONDS) : collection;
+    }
 }