feat: add supports for pagination expression for new api (#98)

* wip * feat: adds support for pagination * Moved test case to DocStoreTest
hypertrace · Jul 13, 2022 · 4df34ac · 4df34ac
1 parent 46ebfdc
commit 4df34ac
Show file tree

Hide file tree

Showing 11 changed files with 275 additions and 234 deletions.
diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreTest.java
@@ -12,8 +12,10 @@
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.GT;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.GTE;
+import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.IN;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.LTE;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NEQ;
+import static org.hypertrace.core.documentstore.expression.operators.SortOrder.ASC;
 import static org.hypertrace.core.documentstore.expression.operators.SortOrder.DESC;
 import static org.hypertrace.core.documentstore.utils.CreateUpdateTestThread.FAILURE;
 import static org.hypertrace.core.documentstore.utils.CreateUpdateTestThread.SUCCESS;
@@ -54,6 +56,11 @@
 import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
 import org.hypertrace.core.documentstore.mongo.MongoDatastore;
 import org.hypertrace.core.documentstore.postgres.PostgresDatastore;
+import org.hypertrace.core.documentstore.query.Pagination;
+import org.hypertrace.core.documentstore.query.Selection;
+import org.hypertrace.core.documentstore.query.SelectionSpec;
+import org.hypertrace.core.documentstore.query.Sort;
+import org.hypertrace.core.documentstore.query.SortingSpec;
 import org.hypertrace.core.documentstore.utils.CreateUpdateTestThread;
 import org.hypertrace.core.documentstore.utils.CreateUpdateTestThread.Operation;
 import org.hypertrace.core.documentstore.utils.Utils;
@@ -1551,12 +1558,22 @@ public void testQueryV1FilterWithNestedFiled(String dataStoreName) throws IOExce
     org.hypertrace.core.documentstore.query.Query query =
         org.hypertrace.core.documentstore.query.Query.builder()
             .setFilter(
-                RelationalExpression.of(
-                    IdentifierExpression.of("quantity"), NEQ, ConstantExpression.of(10)))
+                LogicalExpression.builder()
+                    .operator(AND)
+                    .operand(
+                        RelationalExpression.of(
+                            IdentifierExpression.of("quantity"), GT, ConstantExpression.of(5)))
+                    .operand(
+                        RelationalExpression.of(
+                            IdentifierExpression.of("props.seller.address.city"),
+                            EQ,
+                            ConstantExpression.of("Kolkata")))
+                    .build())
             .build();
 
     Iterator<Document> iterator = collection.aggregate(query);
-    assertSizeAndDocsEqual(dataStoreName, iterator, 6, "mongo/simple_filter_quantity_neq_10.json");
+    assertSizeAndDocsEqual(
+        dataStoreName, iterator, 1, "mongo/test_nest_field_filter_response.json");
   }
 
   @ParameterizedTest
@@ -1772,8 +1789,7 @@ public void testQueryV1AggregationFilterWithWhereClause(String dataStoreName) th
 
   @ParameterizedTest
   @MethodSource("databaseContextProvider")
-  public void testQueryQ1AggregationFilterAlongWithNonAliasFields(String dataStoreName)
-      throws IOException {
+  public void testQueryV1DistinctCountWithSortingSpecs(String dataStoreName) throws IOException {
     Map<Key, Document> documents = createDocumentsFromResource("mongo/collection_data.json");
     Datastore datastore = datastoreMap.get(dataStoreName);
     Collection collection = datastore.getCollection(COLLECTION_NAME);
@@ -1788,53 +1804,64 @@ public void testQueryQ1AggregationFilterAlongWithNonAliasFields(String dataStore
                 AggregateExpression.of(DISTINCT_COUNT, IdentifierExpression.of("quantity")),
                 "qty_count")
             .addSelection(IdentifierExpression.of("item"))
-            .addSelection(IdentifierExpression.of("price"))
             .addAggregation(IdentifierExpression.of("item"))
-            .addAggregation(IdentifierExpression.of("price"))
             .setAggregationFilter(
-                LogicalExpression.builder()
-                    .operator(AND)
-                    .operand(
-                        RelationalExpression.of(
-                            IdentifierExpression.of("qty_count"), LTE, ConstantExpression.of(10)))
-                    .operand(
-                        RelationalExpression.of(
-                            IdentifierExpression.of("price"), GT, ConstantExpression.of(5)))
-                    .build())
+                RelationalExpression.of(
+                    IdentifierExpression.of("qty_count"), LTE, ConstantExpression.of(1000)))
+            .addSort(IdentifierExpression.of("qty_count"), DESC)
+            .addSort(IdentifierExpression.of("item"), DESC)
             .build();
 
     Iterator<Document> resultDocs = collection.aggregate(query);
-    assertSizeAndDocsEqual(
-        dataStoreName, resultDocs, 4, "mongo/test_aggr_alias_distinct_count_response.json");
+    assertDocsAndSizeEqual(resultDocs, "mongo/distinct_count_response.json", 4);
   }
 
   @ParameterizedTest
   @MethodSource("databaseContextProvider")
-  public void testQueryV1DistinctCountWithSortingSpecs(String dataStoreName) throws IOException {
+  public void testFindWithSortingAndPagination(String datastoreName) throws IOException {
     Map<Key, Document> documents = createDocumentsFromResource("mongo/collection_data.json");
-    Datastore datastore = datastoreMap.get(dataStoreName);
+    Datastore datastore = datastoreMap.get(datastoreName);
     Collection collection = datastore.getCollection(COLLECTION_NAME);
 
     // add docs
     boolean result = collection.bulkUpsert(documents);
     Assertions.assertTrue(result);
+    List<SelectionSpec> selectionSpecs =
+        List.of(
+            SelectionSpec.of(IdentifierExpression.of("item")),
+            SelectionSpec.of(IdentifierExpression.of("price")),
+            SelectionSpec.of(IdentifierExpression.of("quantity")),
+            SelectionSpec.of(IdentifierExpression.of("date")));
+    Selection selection = Selection.builder().selectionSpecs(selectionSpecs).build();
+
+    org.hypertrace.core.documentstore.query.Filter filter =
+        org.hypertrace.core.documentstore.query.Filter.builder()
+            .expression(
+                RelationalExpression.of(
+                    IdentifierExpression.of("item"),
+                    IN,
+                    ConstantExpression.ofStrings(List.of("Mirror", "Comb", "Shampoo", "Bottle"))))
+            .build();
+
+    Sort sort =
+        Sort.builder()
+            .sortingSpec(SortingSpec.of(IdentifierExpression.of("quantity"), DESC))
+            .sortingSpec(SortingSpec.of(IdentifierExpression.of("item"), ASC))
+            .build();
+
+    Pagination pagination = Pagination.builder().offset(1).limit(3).build();
 
     org.hypertrace.core.documentstore.query.Query query =
         org.hypertrace.core.documentstore.query.Query.builder()
-            .addSelection(
-                AggregateExpression.of(DISTINCT_COUNT, IdentifierExpression.of("quantity")),
-                "qty_count")
-            .addSelection(IdentifierExpression.of("item"))
-            .addAggregation(IdentifierExpression.of("item"))
-            .setAggregationFilter(
-                RelationalExpression.of(
-                    IdentifierExpression.of("qty_count"), LTE, ConstantExpression.of(1000)))
-            .addSort(IdentifierExpression.of("qty_count"), DESC)
-            .addSort(IdentifierExpression.of("item"), DESC)
+            .setSelection(selection)
+            .setFilter(filter)
+            .setSort(sort)
+            .setPagination(pagination)
             .build();
 
-    Iterator<Document> resultDocs = collection.aggregate(query);
-    assertDocsAndSizeEqual(resultDocs, "mongo/distinct_count_response.json", 4);
+    Iterator<Document> resultDocs = collection.find(query);
+    Utils.assertDocsAndSizeEqual(
+        resultDocs, "mongo/filter_with_sorting_and_pagination_response.json", 3);
   }
 
   @ParameterizedTest

diff --git a/...nTest/java/org/hypertrace/core/documentstore/mongo/MongoQueryExecutorIntegrationTest.java b/...nTest/java/org/hypertrace/core/documentstore/mongo/MongoQueryExecutorIntegrationTest.java
@@ -2,6 +2,7 @@
 
 import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.COUNT;
 import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.DISTINCT;
+import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.DISTINCT_COUNT;
 import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.SUM;
 import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.LENGTH;
 import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.MULTIPLY;
@@ -51,6 +52,7 @@
 import org.hypertrace.core.documentstore.query.SelectionSpec;
 import org.hypertrace.core.documentstore.query.Sort;
 import org.hypertrace.core.documentstore.query.SortingSpec;
+import org.hypertrace.core.documentstore.utils.Utils;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -167,46 +169,6 @@ public void testFindWithDuplicateSelections() throws IOException {
     assertSizeEqual(query, "mongo/simple_filter_response.json");
   }
 
-  @Test
-  public void testFindWithSortingAndPagination() throws IOException {
-    List<SelectionSpec> selectionSpecs =
-        List.of(
-            SelectionSpec.of(IdentifierExpression.of("item")),
-            SelectionSpec.of(IdentifierExpression.of("price")),
-            SelectionSpec.of(IdentifierExpression.of("quantity")),
-            SelectionSpec.of(IdentifierExpression.of("date")));
-    Selection selection = Selection.builder().selectionSpecs(selectionSpecs).build();
-
-    Filter filter =
-        Filter.builder()
-            .expression(
-                RelationalExpression.of(
-                    IdentifierExpression.of("item"),
-                    IN,
-                    ConstantExpression.ofStrings(List.of("Mirror", "Comb", "Shampoo", "Bottle"))))
-            .build();
-
-    Sort sort =
-        Sort.builder()
-            .sortingSpec(SortingSpec.of(IdentifierExpression.of("quantity"), DESC))
-            .sortingSpec(SortingSpec.of(IdentifierExpression.of("item"), ASC))
-            .build();
-
-    Pagination pagination = Pagination.builder().offset(1).limit(3).build();
-
-    Query query =
-        Query.builder()
-            .setSelection(selection)
-            .setFilter(filter)
-            .setSort(sort)
-            .setPagination(pagination)
-            .build();
-
-    Iterator<Document> resultDocs = collection.find(query);
-    assertDocsEqual(resultDocs, "mongo/filter_with_sorting_and_pagination_response.json");
-    assertSizeEqual(query, "mongo/filter_with_sorting_and_pagination_response.json");
-  }
-
   @Test
   public void testFindWithDuplicateSortingAndPagination() throws IOException {
     List<SelectionSpec> selectionSpecs =
@@ -596,6 +558,34 @@ public void testFilterAndUnnest() throws IOException {
     assertSizeEqual(query, "mongo/unwind_filter_response.json");
   }
 
+  @Test
+  public void testQueryQ1AggregationFilterAlongWithNonAliasFields() throws IOException {
+    org.hypertrace.core.documentstore.query.Query query =
+        org.hypertrace.core.documentstore.query.Query.builder()
+            .addSelection(
+                AggregateExpression.of(DISTINCT_COUNT, IdentifierExpression.of("quantity")),
+                "qty_count")
+            .addSelection(IdentifierExpression.of("item"))
+            .addSelection(IdentifierExpression.of("price"))
+            .addAggregation(IdentifierExpression.of("item"))
+            .addAggregation(IdentifierExpression.of("price"))
+            .setAggregationFilter(
+                LogicalExpression.builder()
+                    .operator(AND)
+                    .operand(
+                        RelationalExpression.of(
+                            IdentifierExpression.of("qty_count"), LTE, ConstantExpression.of(10)))
+                    .operand(
+                        RelationalExpression.of(
+                            IdentifierExpression.of("price"), GT, ConstantExpression.of(5)))
+                    .build())
+            .build();
+
+    Iterator<Document> resultDocs = collection.aggregate(query);
+    Utils.assertDocsAndSizeEqualWithoutOrder(
+        resultDocs, "mongo/test_aggr_alias_distinct_count_response.json", 4);
+  }
+
   private static void assertDocsEqual(Iterator<Document> documents, String filePath)
       throws IOException {
     String fileContent = readFileFromResource(filePath).orElseThrow();

diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/utils/Utils.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/utils/Utils.java
@@ -112,4 +112,23 @@ public static void assertDocsAndSizeEqual(
     assertEquals(expected, actual);
     assertEquals(expectedSize, actualSize);
   }
+
+  public static void assertDocsAndSizeEqualWithoutOrder(
+      Iterator<Document> documents, String filePath, int expectedSize) throws IOException {
+    String fileContent = readFileFromResource(filePath).orElseThrow();
+    List<Map<String, Object>> expectedDocs = convertJsonToMap(fileContent);
+
+    List<Map<String, Object>> actualDocs = new ArrayList<>();
+    int actualSize = 0;
+    while (documents.hasNext()) {
+      Map<String, Object> doc = convertDocumentToMap(documents.next());
+      actualDocs.add(doc);
+      actualSize++;
+    }
+
+    long count =
+        expectedDocs.stream().filter(expectedDoc -> actualDocs.contains(expectedDoc)).count();
+    assertEquals(expectedSize, actualSize);
+    assertEquals(expectedSize, count);
+  }
 }
diff --git a/...nt-store/src/integrationTest/resources/mongo/test_aggr_alias_distinct_count_response.json b/...nt-store/src/integrationTest/resources/mongo/test_aggr_alias_distinct_count_response.json
@@ -4,6 +4,11 @@
     "qty_count":1,
     "price":20
   },
+  {
+    "item":"Soap",
+    "qty_count":2,
+    "price":10
+  },
   {
     "item":"Mirror",
     "qty_count":1,
@@ -13,10 +18,5 @@
     "item":"Comb",
     "qty_count":2,
     "price":7.5
-  },
-  {
-    "item":"Soap",
-    "qty_count":2,
-    "price":10
   }
 ]