Enable disk usage for aggregations on MongoDB (#197)

Author: suresh-prakash

.github/workflows/pr-test.yml

@@ -21,23 +21,25 @@ jobs:
           args: jacocoTestReport
 
       - name: Upload coverage to Codecov
-        uses: codecov/codecov-action@v3
+        uses: codecov/codecov-action@v4
         with:
           name: unit test reports
           fail_ci_if_error: true
           flags: unit
+          token: ${{ secrets.CODECOV_TOKEN }}
 
       - name: Integration test
         uses: hypertrace/github-actions/gradle@main
         with:
           args: jacocoIntegrationTestReport
 
       - name: Upload coverage to Codecov
-        uses: codecov/codecov-action@v3
+        uses: codecov/codecov-action@v4
         with:
           name: integration test reports
           fail_ci_if_error: true
           flags: integration
+          token: ${{ secrets.CODECOV_TOKEN }}
 
       - name: copy test reports
         uses: hypertrace/github-actions/gradle@main

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

@@ -130,7 +130,8 @@ public MongoCursor<BasicDBObject> aggregate(final Query originalQuery) {
     logPipeline(pipeline);
 
     try {
-      final AggregateIterable<BasicDBObject> iterable = collection.aggregate(pipeline);
+      final AggregateIterable<BasicDBObject> iterable =
+          collection.aggregate(pipeline).allowDiskUse(true);
       return iterable.cursor();
     } catch (final MongoCommandException e) {
       log.error("Execution failed for query: {}. Aggregation Pipeline: {}", query, pipeline);

document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java

@@ -378,6 +378,7 @@ void testBulkUpdateWithFilter() throws IOException {
 
       when(mockClock.millis()).thenReturn(1660721309000L);
       when(collection.aggregate(anyList())).thenReturn(mockIterable);
+      when(mockIterable.allowDiskUse(true)).thenReturn(mockIterable);
       when(mockIterable.cursor()).thenReturn(mockCursor);
       when(mockCursor.hasNext()).thenReturn(true).thenReturn(false);
       when(mockCursor.next()).thenReturn(response);

document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoQueryExecutorTest.java

@@ -96,6 +96,7 @@ void setUp() {
     when(iterable.sort(any(BasicDBObject.class))).thenReturn(iterable);
 
     when(iterable.cursor()).thenReturn(cursor);
+    when(aggIterable.allowDiskUse(true)).thenReturn(aggIterable);
     when(aggIterable.cursor()).thenReturn(cursor);
   }
 
@@ -522,6 +523,7 @@ private void testAggregation(Query query, final String filePath)
     executor.aggregate(query);
     verify(collection).getNamespace();
     verify(collection).aggregate(pipeline);
+    verify(aggIterable).allowDiskUse(true);
     verify(aggIterable).cursor();
   }