2323import org .springframework .data .mongodb .core .ReadConcernAware ;
2424import org .springframework .data .mongodb .core .ReadPreferenceAware ;
2525import org .springframework .data .mongodb .core .query .Collation ;
26+ import org .springframework .data .mongodb .core .query .DiskUse ;
2627import org .springframework .data .mongodb .util .BsonUtils ;
2728import org .springframework .lang .Contract ;
2829import org .springframework .util .Assert ;
@@ -60,7 +61,7 @@ public class AggregationOptions implements ReadConcernAware, ReadPreferenceAware
6061 private static final String MAX_TIME = "maxTimeMS" ;
6162 private static final String HINT = "hint" ;
6263
63- private final Optional < Boolean > allowDiskUse ;
64+ private final DiskUse diskUse ;
6465 private final boolean explain ;
6566 private final Optional <Document > cursor ;
6667 private final Optional <Collation > collation ;
@@ -85,6 +86,15 @@ public AggregationOptions(boolean allowDiskUse, boolean explain, @Nullable Docum
8586 this (allowDiskUse , explain , cursor , null );
8687 }
8788
89+ public AggregationOptions (DiskUse diskUse , boolean explain , @ Nullable Document cursor ) {
90+ this (diskUse , explain , cursor , null );
91+ }
92+
93+ public AggregationOptions (DiskUse allowDiskUse , boolean explain , @ Nullable Document cursor ,
94+ @ Nullable Collation collation ) {
95+ this (allowDiskUse , explain , cursor , collation , null );
96+ }
97+
8898 /**
8999 * Creates a new {@link AggregationOptions}.
90100 *
@@ -97,7 +107,7 @@ public AggregationOptions(boolean allowDiskUse, boolean explain, @Nullable Docum
97107 */
98108 public AggregationOptions (boolean allowDiskUse , boolean explain , @ Nullable Document cursor ,
99109 @ Nullable Collation collation ) {
100- this (allowDiskUse , explain , cursor , collation , null , null );
110+ this (DiskUse . of ( allowDiskUse ) , explain , cursor , collation , null , null );
101111 }
102112
103113 /**
@@ -113,13 +123,18 @@ public AggregationOptions(boolean allowDiskUse, boolean explain, @Nullable Docum
113123 */
114124 public AggregationOptions (boolean allowDiskUse , boolean explain , @ Nullable Document cursor ,
115125 @ Nullable Collation collation , @ Nullable String comment ) {
126+ this (allowDiskUse ? DiskUse .ALLOW : DiskUse .DENY , explain , cursor , collation , comment , null );
127+ }
128+
129+ public AggregationOptions (DiskUse allowDiskUse , boolean explain , @ Nullable Document cursor ,
130+ @ Nullable Collation collation , @ Nullable String comment ) {
116131 this (allowDiskUse , explain , cursor , collation , comment , null );
117132 }
118133
119134 /**
120135 * Creates a new {@link AggregationOptions}.
121136 *
122- * @param allowDiskUse whether to off-load intensive sort-operations to disk.
137+ * @param diskUse whether to off-load intensive sort-operations to disk.
123138 * @param explain whether to get the execution plan for the aggregation instead of the actual results.
124139 * @param cursor can be {@literal null}, used to pass additional options (such as {@code batchSize}) to the
125140 * aggregation.
@@ -128,10 +143,10 @@ public AggregationOptions(boolean allowDiskUse, boolean explain, @Nullable Docum
128143 * @param hint can be {@literal null}, used to provide an index that would be forcibly used by query optimizer.
129144 * @since 3.1
130145 */
131- private AggregationOptions (@ Nullable Boolean allowDiskUse , boolean explain , @ Nullable Document cursor ,
146+ private AggregationOptions (DiskUse diskUse , boolean explain , @ Nullable Document cursor ,
132147 @ Nullable Collation collation , @ Nullable String comment , @ Nullable Object hint ) {
133148
134- this .allowDiskUse = Optional . ofNullable ( allowDiskUse ) ;
149+ this .diskUse = diskUse ;
135150 this .explain = explain ;
136151 this .cursor = Optional .ofNullable (cursor );
137152 this .collation = Optional .ofNullable (collation );
@@ -172,7 +187,7 @@ public static AggregationOptions fromDocument(Document document) {
172187 String comment = document .getString (COMMENT );
173188 Document hint = document .get (HINT , Document .class );
174189
175- AggregationOptions options = new AggregationOptions (allowDiskUse , explain , cursor , collation , comment , hint );
190+ AggregationOptions options = new AggregationOptions (DiskUse . of ( allowDiskUse ) , explain , cursor , collation , comment , hint );
176191 if (document .containsKey (MAX_TIME )) {
177192 options .maxTime = Duration .ofMillis (document .getLong (MAX_TIME ));
178193 }
@@ -196,7 +211,7 @@ public static Builder builder() {
196211 * @return {@literal true} if enabled; {@literal false} otherwise (or if not set).
197212 */
198213 public boolean isAllowDiskUse () {
199- return allowDiskUse . orElse ( false );
214+ return diskUse . equals ( DiskUse . ALLOW );
200215 }
201216
202217 /**
@@ -206,7 +221,7 @@ public boolean isAllowDiskUse() {
206221 * @since 4.2.5
207222 */
208223 public boolean isAllowDiskUseSet () {
209- return allowDiskUse . isPresent ( );
224+ return ! diskUse . equals ( DiskUse . DEFAULT );
210225 }
211226
212227 /**
@@ -427,7 +442,7 @@ static Document createCursor(int cursorBatchSize) {
427442 */
428443 public static class Builder {
429444
430- private @ Nullable Boolean allowDiskUse ;
445+ private @ Nullable DiskUse diskUse = DiskUse . DEFAULT ;
431446 private boolean explain ;
432447 private @ Nullable Document cursor ;
433448 private @ Nullable Collation collation ;
@@ -447,8 +462,19 @@ public static class Builder {
447462 */
448463 @ Contract ("_ -> this" )
449464 public Builder allowDiskUse (boolean allowDiskUse ) {
465+ return diskUse (DiskUse .of (allowDiskUse ));
466+ }
467+
468+ /**
469+ * Defines whether to off-load intensive sort-operations to disk.
470+ *
471+ * @param diskUse use {@literal true} to allow disk use during the aggregation.
472+ * @return this.
473+ */
474+ @ Contract ("_ -> this" )
475+ public Builder diskUse (DiskUse diskUse ) {
450476
451- this .allowDiskUse = allowDiskUse ;
477+ this .diskUse = diskUse ;
452478 return this ;
453479 }
454480
@@ -655,7 +681,7 @@ public Builder noMapping() {
655681 @ Contract ("-> new" )
656682 public AggregationOptions build () {
657683
658- AggregationOptions options = new AggregationOptions (allowDiskUse , explain , cursor , collation , comment , hint );
684+ AggregationOptions options = new AggregationOptions (diskUse , explain , cursor , collation , comment , hint );
659685 if (maxTime != null ) {
660686 options .maxTime = maxTime ;
661687 }
0 commit comments