Functionality working for MongoDB

hypertrace · Dec 30, 2023 · 0021c9a · 0021c9a
1 parent 7ac51bd
commit 0021c9a
Show file tree

Hide file tree

Showing 10 changed files with 206 additions and 93 deletions.
diff --git a/...ore/src/integrationTest/java/org/hypertrace/core/documentstore/ArrayFiltersQueryTest.java b/...ore/src/integrationTest/java/org/hypertrace/core/documentstore/ArrayFiltersQueryTest.java
@@ -16,7 +16,6 @@
 import java.util.stream.Collector;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
-import org.bson.json.JsonObject;
 import org.hypertrace.core.documentstore.expression.impl.ArrayRelationalFilterExpression;
 import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
 import org.hypertrace.core.documentstore.expression.impl.DocumentArrayFilterExpression;
@@ -31,6 +30,7 @@
 import org.hypertrace.core.documentstore.utils.Utils;
 import org.json.JSONArray;
 import org.json.JSONException;
+import org.json.JSONObject;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.extension.ExtensionContext;
@@ -216,7 +216,14 @@ private String iteratorToJson(final Iterator<Document> iterator) {
     return StreamSupport.stream(
             Spliterators.spliteratorUnknownSize(iterator, Spliterator.ORDERED), false)
         .map(Document::toJson)
-        .map(JsonObject::new)
+        .map(
+            json -> {
+              try {
+                return new JSONObject(json);
+              } catch (JSONException e) {
+                throw new RuntimeException(e);
+              }
+            })
         .collect(
             Collector.of(
                 JSONArray::new, JSONArray::put, (array1, array2) -> array1, JSONArray::toString));

diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoUtils.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoUtils.java
@@ -65,6 +65,11 @@ public static String decodeKey(final String key) {
     return key.replace("\\u002e", FIELD_SEPARATOR).replace("\\u0024", PREFIX).replace("\\\\", "\\");
   }
 
+  public static String getLastField(final String fieldPath) {
+    final String[] fields = fieldPath.split("\\" + FIELD_SEPARATOR);
+    return fields[fields.length - 1];
+  }
+
   public static String sanitizeJsonString(final String jsonString) throws JsonProcessingException {
     final JsonNode jsonNode = MAPPER.readTree(jsonString);
     // escape "." and "$" in field names since Mongo DB does not like them

diff --git a/...rg/hypertrace/core/documentstore/mongo/query/parser/MongoArrayRelationalFilterParser.java b/...rg/hypertrace/core/documentstore/mongo/query/parser/MongoArrayRelationalFilterParser.java
@@ -13,6 +13,7 @@
 import org.hypertrace.core.documentstore.mongo.MongoUtils;
 
 class MongoArrayRelationalFilterParser {
+  private static final String EXPR = "$expr";
   private static final String ANY_ELEMENT_TRUE = "$anyElementTrue";
   private static final String MAP = "$map";
   private static final String INPUT = "input";
@@ -27,10 +28,13 @@ class MongoArrayRelationalFilterParser {
       new MongoIdentifierExpressionParser();
 
   private final UnaryOperator<MongoSelectTypeExpressionParser> wrappingLhsParser;
+  private final boolean exprTypeFilter;
 
   MongoArrayRelationalFilterParser(
-      final UnaryOperator<MongoSelectTypeExpressionParser> wrappingLhsParser) {
+      final UnaryOperator<MongoSelectTypeExpressionParser> wrappingLhsParser,
+      boolean exprTypeFilter) {
     this.wrappingLhsParser = wrappingLhsParser;
+    this.exprTypeFilter = exprTypeFilter;
   }
 
   Map<String, Object> parse(final ArrayRelationalFilterExpression arrayFilterExpression) {
@@ -43,52 +47,60 @@ Map<String, Object> parse(final ArrayRelationalFilterExpression arrayFilterExpre
 
     final SelectTypeExpression lhs = arrayFilterExpression.getFilter().getLhs();
     final String lhsFieldName = lhs.accept(identifierParser);
-    final String alias = MongoUtils.encodeKey(lhsFieldName);
+    final String alias = MongoUtils.getLastField(lhsFieldName);
 
     /*
-     * Wrapping parser to convert 'lhs' to '$prefix.lhs' in the case of nested array filters.
-     * Identifier prefixing parser to convert '$prefix.lhs' to '$$prefix.lhs' in the case of nested array filters.
-     * In the case of non-nested array filters, 'lhs' will just be converted to '$lhs' by the identifier prefixing parser (because the wrapping parser will be identity)
+     * Wrapping parser to convert 'lhs' to '$$prefix.lhs' in the case of nested array filters.
+     * Dollar prefixing idempotent parser to retain '$$prefix.lhs' to '$$prefix.lhs' in the case of nested array filters.
+     * In the case of non-nested array filters, 'lhs' will just be converted to '$lhs' by the dollar prefixing idempotent parser
      */
     final MongoSelectTypeExpressionParser wrappingParser =
-        new MongoIdentifierPrefixingParser(wrappingLhsParser.apply(identifierParser));
+        new MongoDollarPrefixingIdempotentParser(wrappingLhsParser.apply(identifierParser));
     final String mapInput = lhs.accept(wrappingParser);
 
     /*
     {
-      "$anyElementTrue":
-        {
-          "$map":
+      "$expr": {
+        "$anyElementTrue":
           {
-            "input":
+            "$map":
             {
-              "$ifNull": [
-                "$elements",
-                []
-              ]
-            },
-            "as": "elements",
-            "in":
-            {
-              "$eq": ["$$elements", "Water"]
+              "input":
+              {
+                "$ifNull": [
+                  "$elements",
+                  []
+                ]
+              },
+              "as": "elements",
+              "in":
+              {
+                "$eq": ["$$elements", "Water"]
+              }
             }
           }
         }
       }
      */
 
-    return Map.of(
-        operator,
+    final Object filter =
+        arrayFilterExpression
+            .getFilter()
+            .accept(
+                new MongoFilterTypeExpressionParser(
+                    parser -> buildSubstitutingParser(lhsFieldName, alias, parser), true));
+
+    final Map<String, Object> arrayFilter =
         Map.of(
-            MAP,
-            Map.ofEntries(
-                entry(INPUT, Map.of(IF_NULL, new Object[] {mapInput, new Object[0]})),
-                entry(AS, alias),
-                entry(
-                    IN,
-                    new MongoRelationalExpressionParser(
-                            parser -> buildSubstitutingParser(lhsFieldName, alias, parser))
-                        .parse(arrayFilterExpression.getFilter())))));
+            operator,
+            Map.of(
+                MAP,
+                Map.ofEntries(
+                    entry(INPUT, Map.of(IF_NULL, new Object[] {mapInput, new Object[0]})),
+                    entry(AS, alias),
+                    entry(IN, filter))));
+    // If already wrapped inside `$expr` avoid wrapping again
+    return exprTypeFilter ? arrayFilter : Map.of(EXPR, arrayFilter);
   }
 
   private MongoIdentifierPrefixingParser buildSubstitutingParser(

diff --git a/.../org/hypertrace/core/documentstore/mongo/query/parser/MongoDocumentArrayFilterParser.java b/.../org/hypertrace/core/documentstore/mongo/query/parser/MongoDocumentArrayFilterParser.java
@@ -13,6 +13,7 @@
 import org.hypertrace.core.documentstore.mongo.MongoUtils;
 
 class MongoDocumentArrayFilterParser {
+  private static final String EXPR = "$expr";
   private static final String ANY_ELEMENT_TRUE = "$anyElementTrue";
   private static final String MAP = "$map";
   private static final String INPUT = "input";
@@ -27,10 +28,13 @@ class MongoDocumentArrayFilterParser {
       new MongoIdentifierExpressionParser();
 
   private final UnaryOperator<MongoSelectTypeExpressionParser> wrappingLhsParser;
+  private final boolean exprTypeFilter;
 
   MongoDocumentArrayFilterParser(
-      final UnaryOperator<MongoSelectTypeExpressionParser> wrappingLhsParser) {
+      final UnaryOperator<MongoSelectTypeExpressionParser> wrappingLhsParser,
+      final boolean exprTypeFilter) {
     this.wrappingLhsParser = wrappingLhsParser;
+    this.exprTypeFilter = exprTypeFilter;
   }
 
   Map<String, Object> parse(final DocumentArrayFilterExpression arrayFilterExpression) {
@@ -43,54 +47,60 @@ Map<String, Object> parse(final DocumentArrayFilterExpression arrayFilterExpress
 
     final SelectTypeExpression arraySource = arrayFilterExpression.getArraySource();
     final String arraySourceName = arraySource.accept(identifierParser);
-    final String alias = MongoUtils.encodeKey(arraySourceName);
+    final String alias = MongoUtils.getLastField(arraySourceName);
 
     /*
-     * Wrapping parser to convert 'lhs' to '$prefix.lhs' in the case of nested array filters.
-     * Identifier prefixing parser to convert '$prefix.lhs' to '$$prefix.lhs' in the case of nested array filters.
-     * In the case of non-nested array filters, 'lhs' will just be converted to '$lhs' by the identifier prefixing parser (because the wrapping parser will be identity)
+     * Wrapping parser to convert 'lhs' to '$$prefix.lhs' in the case of nested array filters.
+     * Dollar prefixing idempotent parser to retain '$$prefix.lhs' to '$$prefix.lhs' in the case of nested array filters.
+     * In the case of non-nested array filters, 'lhs' will just be converted to '$lhs' by the dollar prefixing idempotent parser
      */
     final MongoSelectTypeExpressionParser wrappingParser =
-        new MongoIdentifierPrefixingParser(wrappingLhsParser.apply(identifierParser));
+        new MongoDollarPrefixingIdempotentParser(wrappingLhsParser.apply(identifierParser));
     final String mapInput = arraySource.accept(wrappingParser);
 
     /*
     {
-      "$anyElementTrue":
-        {
-          "$map":
+      "$expr": {
+        "$anyElementTrue":
           {
-            "input":
+            "$map":
             {
-              "$ifNull": [
-                "$planets",
-                []
-              ]
-            },
-            "as": "planets",
-            "in":
-            {
-              "$eq": ["$$planets.name", "Mars"]
+              "input":
+              {
+                "$ifNull": [
+                  "$planets",
+                  []
+                ]
+              },
+              "as": "planets",
+              "in":
+              {
+                "$eq": ["$$planets.name", "Mars"]
+              }
             }
           }
         }
       }
      */
 
-    return Map.of(
-        operator,
+    final Object filter =
+        arrayFilterExpression
+            .getFilter()
+            .accept(
+                new MongoFilterTypeExpressionParser(
+                    parser -> buildPrefixingParser(alias, parser), true));
+
+    final Map<String, Object> arrayFilter =
         Map.of(
-            MAP,
-            Map.ofEntries(
-                entry(INPUT, Map.of(IF_NULL, new Object[] {mapInput, new Object[0]})),
-                entry(AS, alias),
-                entry(
-                    IN,
-                    arrayFilterExpression
-                        .getFilter()
-                        .accept(
-                            new MongoFilterTypeExpressionParser(
-                                parser -> buildPrefixingParser(alias, parser)))))));
+            operator,
+            Map.of(
+                MAP,
+                Map.ofEntries(
+                    entry(INPUT, Map.of(IF_NULL, new Object[] {mapInput, new Object[0]})),
+                    entry(AS, alias),
+                    entry(IN, filter))));
+    // If already wrapped inside `$expr` avoid wrapping again
+    return exprTypeFilter ? arrayFilter : Map.of(EXPR, arrayFilter);
   }
 
   private MongoIdentifierPrefixingParser buildPrefixingParser(

diff --git a/...ypertrace/core/documentstore/mongo/query/parser/MongoDollarPrefixingIdempotentParser.java b/...ypertrace/core/documentstore/mongo/query/parser/MongoDollarPrefixingIdempotentParser.java
@@ -0,0 +1,25 @@
+package org.hypertrace.core.documentstore.mongo.query.parser;
+
+import static org.hypertrace.core.documentstore.mongo.MongoUtils.PREFIX;
+
+import java.util.Optional;
+import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
+
+final class MongoDollarPrefixingIdempotentParser extends MongoSelectTypeExpressionParser {
+  MongoDollarPrefixingIdempotentParser(final MongoSelectTypeExpressionParser baseParser) {
+    super(baseParser);
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public String visit(final IdentifierExpression expression) {
+    return Optional.ofNullable(baseParser.visit(expression))
+        .map(Object::toString)
+        .map(this::idempotentPrefix)
+        .orElse(null);
+  }
+
+  private String idempotentPrefix(final String identifier) {
+    return identifier.startsWith(PREFIX) ? identifier : PREFIX + identifier;
+  }
+}
diff --git a/.../hypertrace/core/documentstore/mongo/query/parser/MongoExprRelationalFilterOperation.java b/.../hypertrace/core/documentstore/mongo/query/parser/MongoExprRelationalFilterOperation.java
@@ -0,0 +1,21 @@
+package org.hypertrace.core.documentstore.mongo.query.parser;
+
+import java.util.Map;
+import lombok.AllArgsConstructor;
+import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
+
+@AllArgsConstructor
+class MongoExprRelationalFilterOperation implements RelationalFilterOperation {
+  public static final String EXPR = "$expr";
+
+  private final MongoSelectTypeExpressionParser lhsParser;
+  private final MongoSelectTypeExpressionParser rhsParser;
+  private final String operator;
+
+  @Override
+  public Map<String, Object> apply(final SelectTypeExpression lhs, final SelectTypeExpression rhs) {
+    final Object parsedLhs = lhs.accept(lhsParser);
+    final Object parsedRhs = rhs.accept(rhsParser);
+    return Map.of(operator, new Object[] {parsedLhs, parsedRhs});
+  }
+}
diff --git a/...org/hypertrace/core/documentstore/mongo/query/parser/MongoFilterTypeExpressionParser.java b/...org/hypertrace/core/documentstore/mongo/query/parser/MongoFilterTypeExpressionParser.java
@@ -22,26 +22,29 @@ public final class MongoFilterTypeExpressionParser implements FilterTypeExpressi
   private static final String FILTER_CLAUSE = "$match";
 
   private final UnaryOperator<MongoSelectTypeExpressionParser> wrappingLhsParser;
+  private final boolean exprTypeFilter;
 
   public MongoFilterTypeExpressionParser() {
-    this(UnaryOperator.identity());
+    this(UnaryOperator.identity(), false);
   }
 
   public MongoFilterTypeExpressionParser(
-      final UnaryOperator<MongoSelectTypeExpressionParser> wrappingLhsParser) {
+      final UnaryOperator<MongoSelectTypeExpressionParser> wrappingLhsParser,
+      final boolean exprTypeFilter) {
     this.wrappingLhsParser = wrappingLhsParser;
+    this.exprTypeFilter = exprTypeFilter;
   }
 
   @SuppressWarnings("unchecked")
   @Override
   public Map<String, Object> visit(final LogicalExpression expression) {
-    return new MongoLogicalExpressionParser(wrappingLhsParser).parse(expression);
+    return new MongoLogicalExpressionParser(wrappingLhsParser, exprTypeFilter).parse(expression);
   }
 
   @SuppressWarnings("unchecked")
   @Override
   public Map<String, Object> visit(final RelationalExpression expression) {
-    return new MongoRelationalExpressionParser(wrappingLhsParser).parse(expression);
+    return new MongoRelationalExpressionParser(wrappingLhsParser, exprTypeFilter).parse(expression);
   }
 
   @SuppressWarnings("unchecked")
@@ -57,13 +60,14 @@ public Map<String, Object> visit(final KeyExpression expression) {
   @SuppressWarnings("unchecked")
   @Override
   public Map<String, Object> visit(final ArrayRelationalFilterExpression expression) {
-    return new MongoArrayRelationalFilterParser(wrappingLhsParser).parse(expression);
+    return new MongoArrayRelationalFilterParser(wrappingLhsParser, exprTypeFilter)
+        .parse(expression);
   }
 
   @SuppressWarnings("unchecked")
   @Override
   public Map<String, Object> visit(final DocumentArrayFilterExpression expression) {
-    return new MongoDocumentArrayFilterParser(wrappingLhsParser).parse(expression);
+    return new MongoDocumentArrayFilterParser(wrappingLhsParser, exprTypeFilter).parse(expression);
   }
 
   public static BasicDBObject getFilterClause(

diff --git a/...ertrace/core/documentstore/mongo/query/parser/MongoFunctionRelationalFilterOperation.java b/...ertrace/core/documentstore/mongo/query/parser/MongoFunctionRelationalFilterOperation.java
@@ -1,14 +1,14 @@
 package org.hypertrace.core.documentstore.mongo.query.parser;
 
+import static org.hypertrace.core.documentstore.mongo.query.parser.MongoExprRelationalFilterOperation.EXPR;
+
 import com.mongodb.BasicDBObject;
 import java.util.Map;
 import lombok.AllArgsConstructor;
 import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
 
 @AllArgsConstructor
 class MongoFunctionRelationalFilterOperation implements RelationalFilterOperation {
-  private static final String EXPR = "$expr";
-
   private final MongoSelectTypeExpressionParser functionParser =
       new MongoFunctionExpressionParser();