Handle the zero limit case explicitly for MongoDB aggregation (#187)

Author: suresh-prakash

document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java

@@ -2053,6 +2053,17 @@ public void testAggregateWithSingleKey(final String datastoreName) throws IOExce
     }
   }
 
+  @ParameterizedTest
+  @ArgumentsSource(AllProvider.class)
+  public void testAggregateWithZeroLimitAndOffset(final String datastoreName) throws IOException {
+    final Collection collection = getCollection(datastoreName);
+
+    final Iterator<Document> resultDocs =
+        collection.aggregate(
+            Query.builder().setPagination(Pagination.builder().limit(0).offset(0).build()).build());
+    assertDocsAndSizeEqual(datastoreName, resultDocs, "query/empty_response.json", 0);
+  }
+
   @Nested
   class AtomicUpdateTest {
     @ParameterizedTest

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/MongoPaginationHelper.java

@@ -12,14 +12,16 @@ public class MongoPaginationHelper {
   private static final String LIMIT_CLAUSE = "$limit";
 
   static BasicDBObject getSkipClause(final Query query) {
-    Optional<Pagination> paginationOptional = query.getPagination();
+    Optional<Pagination> paginationOptional =
+        query.getPagination().filter(pagination -> pagination.getOffset() > 0);
     return paginationOptional
         .map(pagination -> new BasicDBObject(SKIP_CLAUSE, pagination.getOffset()))
         .orElse(new BasicDBObject());
   }
 
   static BasicDBObject getLimitClause(final Query query) {
-    Optional<Pagination> paginationOptional = query.getPagination();
+    Optional<Pagination> paginationOptional =
+        query.getPagination().filter(pagination -> pagination.getLimit() > 0);
     return paginationOptional
         .map(pagination -> new BasicDBObject(LIMIT_CLAUSE, pagination.getLimit()))
         .orElse(new BasicDBObject());

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/MongoQueryExecutor.java

@@ -18,11 +18,14 @@
 
 import com.mongodb.BasicDBObject;
 import com.mongodb.MongoCommandException;
+import com.mongodb.ServerAddress;
+import com.mongodb.ServerCursor;
 import com.mongodb.client.AggregateIterable;
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.MongoCursor;
 import java.util.Collection;
 import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -49,6 +52,47 @@ public class MongoQueryExecutor {
               query -> singleton(getSkipClause(query)),
               query -> singleton(getLimitClause(query)));
 
+  private static final Integer ZERO = Integer.valueOf(0);
+  private static final MongoCursor<BasicDBObject> EMPTY_CURSOR =
+      new MongoCursor<>() {
+        @Override
+        public void close() {
+          // Do nothing
+        }
+
+        @Override
+        public boolean hasNext() {
+          return false;
+        }
+
+        @Override
+        public BasicDBObject next() {
+          throw new NoSuchElementException();
+        }
+
+        @Override
+        public int available() {
+          return 0;
+        }
+
+        @Override
+        public BasicDBObject tryNext() {
+          throw new NoSuchElementException();
+        }
+
+        @Override
+        public ServerCursor getServerCursor() {
+          // It is okay to throw exception since we are never invoking this method
+          throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public ServerAddress getServerAddress() {
+          // It is okay to throw exception since we are never invoking this method
+          throw new UnsupportedOperationException();
+        }
+      };
+
   private final com.mongodb.client.MongoCollection<BasicDBObject> collection;
 
   public MongoCursor<BasicDBObject> find(final Query query) {
@@ -70,6 +114,11 @@ public MongoCursor<BasicDBObject> find(final Query query) {
   }
 
   public MongoCursor<BasicDBObject> aggregate(final Query originalQuery) {
+    if (originalQuery.getPagination().map(Pagination::getLimit).map(ZERO::equals).orElse(false)) {
+      log.debug("Not executing query because of a 0 limit");
+      return EMPTY_CURSOR;
+    }
+
     Query query = transformAndLog(originalQuery);
 
     List<BasicDBObject> pipeline =

document-store/src/main/java/org/hypertrace/core/documentstore/query/Pagination.java

@@ -17,6 +17,10 @@ public static class PaginationBuilder {
     public Pagination build() {
       Preconditions.checkArgument(limit != null, "limit is null");
       Preconditions.checkArgument(offset != null, "offset is null");
+
+      Preconditions.checkArgument(limit >= 0, "limit must be non-negative");
+      Preconditions.checkArgument(offset >= 0, "offset must be non-negative");
+
       return new Pagination(limit, offset);
     }
   }

document-store/src/test/resources/mongo/pipeline/with_pagination.json

@@ -6,9 +6,6 @@
       }
     }
   },
-  {
-    "$skip": 0
-  },
   {
     "$limit": 10
   }