Merge pull request #94 from steleal/gh-93

Add Query.reqTotal
Restream · Aug 19, 2022 · 2a6053e · 2a6053e
2 parents 43621b8 + 1d5c27b
commit 2a6053e
Show file tree

Hide file tree

Showing 10 changed files with 99 additions and 56 deletions.
diff --git a/src/main/java/ru/rt/restream/reindexer/Namespace.java b/src/main/java/ru/rt/restream/reindexer/Namespace.java
@@ -111,12 +111,12 @@ public interface Namespace<T> {
     String getMeta(String key);
 
     /**
-     * Executes the given SQL query and returns a {@link CloseableIterator}.
+     * Executes the given SQL query and returns a {@link ResultIterator}.
      *
      * @param query the SQL query to execute
-     * @return the {@link CloseableIterator} to use
+     * @return the {@link ResultIterator} to use
      */
-    CloseableIterator<T> execSql(String query);
+    ResultIterator<T> execSql(String query);
 
     /**
      * Executes the given SQL update query.

diff --git a/src/main/java/ru/rt/restream/reindexer/Query.java b/src/main/java/ru/rt/restream/reindexer/Query.java
@@ -47,6 +47,7 @@
 import static ru.rt.restream.reindexer.binding.Consts.INNER_JOIN;
 import static ru.rt.restream.reindexer.binding.Consts.LEFT_JOIN;
 import static ru.rt.restream.reindexer.binding.Consts.MERGE;
+import static ru.rt.restream.reindexer.binding.Consts.MODE_ACCURATE_TOTAL;
 import static ru.rt.restream.reindexer.binding.Consts.OR_INNER_JOIN;
 import static ru.rt.restream.reindexer.binding.Consts.VALUE_BOOL;
 import static ru.rt.restream.reindexer.binding.Consts.VALUE_NULL;
@@ -611,6 +612,19 @@ public AggregationFacetRequest sort(String field, boolean desc) {
 
     }
 
+
+    /**
+     * Calculate the total count of matching documents.
+     *
+     * @return the {@link Query} for further customizations
+     */
+    public Query<T> reqTotal() {
+        logBuilder.reqTotal();
+        buffer.putVarUInt32(QUERY_REQ_TOTAL)
+                .putVarUInt32(MODE_ACCURATE_TOTAL);
+        return this;
+    }
+
     /**
      * Limit is used to retrieve records from the namespace in a database and limit the number of items returned based
      * on a limit value
@@ -736,7 +750,7 @@ private void putValue(Object value) {
      * @return stream of items
      */
     public Stream<T> stream() {
-        CloseableIterator<T> iterator = execute();
+        ResultIterator<T> iterator = execute();
         Spliterator<T> spliterator = Spliterators.spliterator(iterator, iterator.size(), Spliterator.NONNULL);
         return StreamSupport.stream(spliterator, false).onClose(iterator::close);
     }
@@ -747,7 +761,7 @@ public Stream<T> stream() {
      * @return list of items
      */
     public List<T> toList() {
-        try (CloseableIterator<T> iterator = execute()) {
+        try (ResultIterator<T> iterator = execute()) {
             List<T> result = new ArrayList<>();
             while (iterator.hasNext()) {
                 result.add(iterator.next());
@@ -784,7 +798,7 @@ public Optional<T> findOne() {
     }
 
     private T getOneInternal() {
-        try (CloseableIterator<T> iterator = execute()) {
+        try (ResultIterator<T> iterator = execute()) {
             T item = null;
             if (iterator.hasNext()) {
                 item = iterator.next();
@@ -804,7 +818,7 @@ private T getOneInternal() {
      * @return count of items
      */
     public long count() {
-        try (CloseableIterator<T> iterator = execute()) {
+        try (ResultIterator<T> iterator = execute()) {
             return iterator.size();
         }
     }
@@ -824,7 +838,7 @@ public boolean notExists() {
      * @return true if the item exists
      */
     public boolean exists() {
-        try (CloseableIterator<T> iterator = execute()) {
+        try (ResultIterator<T> iterator = execute()) {
             return iterator.hasNext();
         }
     }
@@ -834,7 +848,7 @@ public boolean exists() {
      *
      * @return an iterator over a query result
      */
-    public CloseableIterator<T> execute() {
+    public ResultIterator<T> execute() {
         return execute(namespace.getItemClass());
     }
 
@@ -845,7 +859,7 @@ public CloseableIterator<T> execute() {
      * @param itemClass the item class
      * @return an iterator over a query result
      */
-    public <S> CloseableIterator<S> execute(Class<S> itemClass) {
+    public <S> ResultIterator<S> execute(Class<S> itemClass) {
         long[] ptVersions = prepareQueryAndGetPayloadTypesVersions();
 
         RequestContext requestContext = transactionContext != null

diff --git a/src/main/java/ru/rt/restream/reindexer/QueryLogBuilder.java b/src/main/java/ru/rt/restream/reindexer/QueryLogBuilder.java
@@ -34,6 +34,7 @@ class QueryLogBuilder {
 
     private String namespace;
     private QueryType type = QueryType.SELECT;
+    private boolean reqTotal;
     private Integer offset;
     private Integer limit;
     private final Map<Query<?>.AggregationFacetRequest, AggregateEntry> facetParams = new HashMap<>();
@@ -140,6 +141,9 @@ String getSql() {
                 stringBuilder.append(" ")
                         .append(getSelectPart());
             }
+            if (reqTotal) {
+                stringBuilder.append(", COUNT(*)");
+            }
             stringBuilder.append(" FROM");
         }
 
@@ -416,6 +420,13 @@ void facetSort(Query<?>.AggregationFacetRequest facet, String sortIndex, boolean
         facetParams.get(facet).sortEntries.add(sortEntry);
     }
 
+    /**
+     * Set flag of request of total count of items.
+     */
+    void reqTotal() {
+        this.reqTotal = true;
+    }
+
     /**
      * Add query offset to builder.
      *

diff --git a/src/main/java/ru/rt/restream/reindexer/QueryResultIterator.java b/src/main/java/ru/rt/restream/reindexer/QueryResultIterator.java
@@ -39,7 +39,7 @@
  * An iterator over a query result.
  * Maintains a cursor pointing to its current row of data. Initially the cursor is positioned before the first row.
  */
-public class QueryResultIterator<T> implements CloseableIterator<T> {
+public class QueryResultIterator<T> implements ResultIterator<T> {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(QueryResultIterator.class);
 
@@ -269,6 +269,11 @@ private void fetchResults() {
         parseQueryResult(queryResult);
     }
 
+    @Override
+    public long getTotalCount() {
+        return queryResult.getTotalCount();
+    }
+
     @Override
     public long size() {
         return queryResult.getQCount();

diff --git a/src/main/java/ru/rt/restream/reindexer/QueryResultJsonIterator.java b/src/main/java/ru/rt/restream/reindexer/QueryResultJsonIterator.java
@@ -22,7 +22,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 
-public class QueryResultJsonIterator implements CloseableIterator<String> {
+public class QueryResultJsonIterator implements ResultIterator<String> {
 
     private final RequestContext requestContext;
 
@@ -113,6 +113,11 @@ public String fetchAll(String rootName) {
         return builder.toString();
     }
 
+    @Override
+    public long getTotalCount() {
+        return queryResult.getTotalCount();
+    }
+
     @Override
     public long size() {
         return queryResult.getQCount();

diff --git a/src/main/java/ru/rt/restream/reindexer/ReindexerNamespace.java b/src/main/java/ru/rt/restream/reindexer/ReindexerNamespace.java
@@ -250,7 +250,7 @@ public String getMeta(String key) {
     }
 
     @Override
-    public CloseableIterator<T> execSql(String query) {
+    public ResultIterator<T> execSql(String query) {
         return reindexer.execSql(query, itemClass);
     }
 

diff --git a/...restream/reindexer/CloseableIterator.java → ...rt/restream/reindexer/ResultIterator.java b/...restream/reindexer/CloseableIterator.java → ...rt/restream/reindexer/ResultIterator.java
@@ -25,7 +25,12 @@
  * @see Iterator
  * @see AutoCloseable
  */
-public interface CloseableIterator<E> extends Iterator<E>, AutoCloseable {
+public interface ResultIterator<E> extends Iterator<E>, AutoCloseable {
+
+    /**
+     * Return the total count of matching documents, if {@link Query} has set reqTotal() flag, else return 0.
+     */
+    long getTotalCount();
 
     /**
      * Returns the iterator size.

diff --git a/src/test/java/ru/rt/restream/reindexer/connector/AggregationTest.java b/src/test/java/ru/rt/restream/reindexer/connector/AggregationTest.java
@@ -17,7 +17,7 @@
 
 import org.junit.jupiter.api.Test;
 import ru.rt.restream.reindexer.AggregationResult;
-import ru.rt.restream.reindexer.CloseableIterator;
+import ru.rt.restream.reindexer.ResultIterator;
 import ru.rt.restream.reindexer.Namespace;
 import ru.rt.restream.reindexer.NamespaceOptions;
 import ru.rt.restream.reindexer.Query;
@@ -48,7 +48,7 @@ public void testDistinct() {
             itemNamespace.insert(item);
         }
 
-        CloseableIterator<Item> result = itemNamespace.query()
+        ResultIterator<Item> result = itemNamespace.query()
                 .aggregateDistinct("name")
                 .execute();
         AggregationResult aggResult = result.aggResults().get(0);
@@ -72,7 +72,7 @@ public void testSum() {
 
         Query<Item> query = itemNamespace.query()
                 .aggregateSum("price");
-        CloseableIterator<Item> result = query.execute();
+        ResultIterator<Item> result = query.execute();
         AggregationResult sumResult = result.aggResults().get(0);
         assertThat(sumResult.getValue(), is(499500D));
     }
@@ -91,7 +91,7 @@ public void testAvg() {
         Query<Item> query = itemNamespace.query()
                 .aggregateAvg("price");
         query.aggregateFacet("id", "price").sort("id", true).offset(10).limit(100);
-        CloseableIterator<Item> result = query.execute();
+        ResultIterator<Item> result = query.execute();
         AggregationResult avgResult = result.aggResults().get(0);
         assertThat(avgResult.getValue(), is(499.5D));
     }
@@ -109,7 +109,7 @@ public void testMax() {
 
         Query<Item> query = itemNamespace.query()
                 .aggregateMax("price");
-        CloseableIterator<Item> result = query.execute();
+        ResultIterator<Item> result = query.execute();
         AggregationResult maxResult = result.aggResults().get(0);
         assertThat(maxResult.getValue(), is(999D));
     }
@@ -127,7 +127,7 @@ public void testMin() {
 
         Query<Item> query = itemNamespace.query()
                 .aggregateMin("price");
-        CloseableIterator<Item> result = query.execute();
+        ResultIterator<Item> result = query.execute();
         AggregationResult minResult = result.aggResults().get(0);
         assertThat(minResult.getValue(), is(0D));
     }
@@ -146,7 +146,7 @@ public void testFacet() {
 
         Query<Item> query = itemNamespace.query();
         query.aggregateFacet("name", "price").sort("name", true).sort("price", false).limit(100);
-        CloseableIterator<Item> result = query.execute();
+        ResultIterator<Item> result = query.execute();
         AggregationResult aggResult = result.aggResults().get(0);
 
         List<AggregationResult.Facet> facets = aggResult.getFacets();