From 98ea5e5ae028210cc1b897cfbbad7b8915c093b4 Mon Sep 17 00:00:00 2001
From: Pavels Fjodorovs <pavels.fjodorovs@hotelplan.com>
Date: Fri, 23 Aug 2024 16:59:56 +0200
Subject: [PATCH] Allow to specify WITHCOUNT in QueryOptions /
 AggregationOptions fixes #6119

Signed-off-by: Pavels Fjodorovs <pavels.fjodorovs@hotelplan.com>
---
 .../main/java/org/redisson/RedissonSearch.java |  6 ++++++
 .../search/aggregate/AggregationOptions.java   | 18 ++++++++++++++++++
 .../api/search/query/QueryOptions.java         | 10 ++++++++++
 3 files changed, 34 insertions(+)

diff --git a/redisson/src/main/java/org/redisson/RedissonSearch.java b/redisson/src/main/java/org/redisson/RedissonSearch.java
index a43da3a74..c4beaf3eb 100644
--- a/redisson/src/main/java/org/redisson/RedissonSearch.java
+++ b/redisson/src/main/java/org/redisson/RedissonSearch.java
@@ -476,6 +476,9 @@ public class RedissonSearch implements RSearch {
             if (options.getSortOrder() != null) {
                 args.add(options.getSortOrder());
             }
+            if (options.isWithCount()) {
+                args.add("WITHCOUNT");
+            }
         }
         if (options.getOffset() != null
                 && options.getCount() != null) {
@@ -588,6 +591,9 @@ public class RedissonSearch implements RSearch {
                 args.add("MAX");
                 args.add(options.getSortedByMax());
             }
+            if (options.isSortedByWithCount()) {
+                args.add("WITHCOUNT");
+            }
         }
         for (Expression expression : options.getExpressions()) {
             args.add("APPLY");
diff --git a/redisson/src/main/java/org/redisson/api/search/aggregate/AggregationOptions.java b/redisson/src/main/java/org/redisson/api/search/aggregate/AggregationOptions.java
index 8c16740d9..25e383637 100644
--- a/redisson/src/main/java/org/redisson/api/search/aggregate/AggregationOptions.java
+++ b/redisson/src/main/java/org/redisson/api/search/aggregate/AggregationOptions.java
@@ -35,6 +35,7 @@ public final class AggregationOptions {
     private List<GroupParams> groupByParams = Collections.emptyList();
     private List<SortedField> sortedByFields = Collections.emptyList();
     private Integer sortedByMax;
+    private boolean sortedByWithCount;
     private List<Expression> expressions = Collections.emptyList();
     private Integer offset;
     private Integer count;
@@ -88,6 +89,19 @@ public final class AggregationOptions {
         return this;
     }
 
+    public AggregationOptions sortBy(boolean withCount, SortedField... fields) {
+        sortedByWithCount = withCount;
+        sortedByFields = Arrays.asList(fields);
+        return this;
+    }
+
+    public AggregationOptions sortBy(int max, boolean withCount, SortedField... fields) {
+        sortedByMax = max;
+        sortedByWithCount = withCount;
+        sortedByFields = Arrays.asList(fields);
+        return this;
+    }
+
     public AggregationOptions apply(Expression... expressions) {
         this.expressions = Arrays.asList(expressions);
         return this;
@@ -160,6 +174,10 @@ public final class AggregationOptions {
         return sortedByMax;
     }
 
+    public boolean isSortedByWithCount() {
+        return sortedByWithCount;
+    }
+
     public List<Expression> getExpressions() {
         return expressions;
     }
diff --git a/redisson/src/main/java/org/redisson/api/search/query/QueryOptions.java b/redisson/src/main/java/org/redisson/api/search/query/QueryOptions.java
index bc0c64266..fcd4d0ef6 100644
--- a/redisson/src/main/java/org/redisson/api/search/query/QueryOptions.java
+++ b/redisson/src/main/java/org/redisson/api/search/query/QueryOptions.java
@@ -44,6 +44,7 @@ public final class QueryOptions {
     private boolean explainScore;
     private String sortBy;
     private SortOrder sortOrder;
+    private boolean withCount;
     private Integer offset;
     private Integer count;
     private Map<String, Object> params = Collections.emptyMap();
@@ -137,6 +138,11 @@ public final class QueryOptions {
         return this;
     }
 
+    public QueryOptions withCount(boolean withCount) {
+        this.withCount = withCount;
+        return this;
+    }
+
     public QueryOptions limit(int offset, int count) {
         this.offset = offset;
         this.count = count;
@@ -234,6 +240,10 @@ public final class QueryOptions {
         return sortBy;
     }
 
+    public boolean isWithCount() {
+        return withCount;
+    }
+
     public Integer getOffset() {
         return offset;
     }