[AD-627] Add new method to get aggregate operations as strings (#295)

* [AD-627] Add new getter method in DocumentDbMqlQueryContext

* Some formatting changes

* Commit Code Coverage Badge

* Rename test file

* Commit Code Coverage Badge

* Improvements from code review - move lambda to one line

* Commit Code Coverage Badge

* [AD-627] Add additional test cases

* [AD-627] Add additional test cases

* Commit Code Coverage Badge

Co-authored-by: andiem-bq <[email protected]>

calcite-adapter/src/main/java/software/amazon/documentdb/jdbc/query/DocumentDbMqlQueryContext.java

@@ -19,12 +19,14 @@
 import lombok.Builder;
 import lombok.Getter;
 import org.bson.conversions.Bson;
+import org.bson.json.JsonMode;
+import org.bson.json.JsonWriterSettings;
 import software.amazon.documentdb.jdbc.common.utilities.JdbcColumnMetaData;
 import java.util.List;
+import java.util.stream.Collectors;
 
 /**
- * This is meant to carry
- * all the information needed to execute the query in MongoDb and
+ * This is meant to carry all the information needed to execute the query in DocumentDB and
  * construct a ResultSet.
  */
 @Getter
@@ -38,4 +40,16 @@ public class DocumentDbMqlQueryContext {
     private final String collectionName;
     /** The path information for the output documents. Maps column names to field paths.*/
     private final List<String> paths;
+
+    /**
+     * Gets the aggregation operations (stages) for the query as a list of strings.
+     *
+     * @return the aggregation operations as an ordered list of strings in extended JSON format.
+     */
+    public List<String> getAggregateOperationsAsStrings() {
+        return aggregateOperations.stream()
+                .map(doc ->
+                        doc.toBsonDocument().toJson(JsonWriterSettings.builder().outputMode(JsonMode.EXTENDED).build()))
+                .collect(Collectors.toList());
+    }
 }

calcite-adapter/src/test/java/software/amazon/documentdb/jdbc/query/DocumentDbMqlQueryContextTest.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright <2021> Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ *
+ */
+
+package software.amazon.documentdb.jdbc.query;
+
+import org.bson.BsonDocument;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class DocumentDbMqlQueryContextTest {
+
+    @Test
+    @DisplayName("Tests that aggregate operations are correctly converted from BSON document to string in getter.")
+    void testGetAggregateOperationsAsStrings() {
+        final List<String> stages = new ArrayList<>();
+        // Generic stage
+        stages.add(
+                "{\"$unwind\": {"
+                        + "\"path\": \"$array\", "
+                        + "\"includeArrayIndex\": \"array_index_lvl_0\", "
+                        + "\"preserveNullAndEmptyArrays\": true}}");
+        // Stage with 3-valued logic (many null checks)
+        stages.add(
+                "{\"$project\": {"
+                        + "\"booleanField\": "
+                        + "{\"$cond\": [{\"$and\": [{\"$gt\": [\"$array.field\", null]}, "
+                        + "{\"$gt\": [\"$array.field2\", null]}]}, "
+                        + "{\"$eq\": [\"$array.field\", \"$array.field2\"]}, null]}}}");
+        // Stage with different Bson types
+        stages.add(
+                "{\"$project\": {"
+                        + "\"literalNull\": {\"$literal\": null}, "
+                        + "\"literalTimestamp\": {\"$date\": {\"$numberLong\": \"1505938660000\"}}, "
+                        + "\"literalInt\": {\"$literal\": {\"$numberInt\": \"-2147483648\"}}, "
+                        + "\"literalDecimal\": {\"$literal\": {\"$numberDouble\": \"123.45\"}}, "
+                        + "\"literalVarchar\": {\"$literal\": \"Hello! 你好!\"}, "
+                        + "\"literalBinary\": {\"$binary\": {\"base64\": \"RfCr\", \"subType\": \"00\"}}}}");
+        // Stage with a lot of nesting and aggregate operators
+        stages.add(
+                "{\"$project\": {\"EXPR$0\": {\"$substrCP\": [\"$array.field\", "
+                        + "{\"$subtract\": [\"$array.field2\", {\"$numberInt\": \"1\"}]}, "
+                        + "{\"$subtract\": [\"$array.field1\", \"$array.field2\"]}]}, "
+                        + "\"_id\": {\"$numberInt\": \"0\"}}}");
+
+        final DocumentDbMqlQueryContext context =
+                DocumentDbMqlQueryContext.builder()
+                        .aggregateOperations(
+                                stages.stream().map(BsonDocument::parse).collect(Collectors.toList()))
+                        .build();
+        Assertions.assertEquals(stages.size(), context.getAggregateOperationsAsStrings().size());
+        for (int i = 0; i < stages.size(); i++) {
+            Assertions.assertEquals(stages.get(i), context.getAggregateOperationsAsStrings().get(i));
+        }
+    }
+}