feat(query-generation): support for addFields INTELLIJ-194 (#129)

mongodb · Jan 23, 2025 · d7678ff · d7678ff
1 parent 99f18ef
commit d7678ff
Show file tree

Hide file tree

Showing 6 changed files with 148 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@ MongoDB plugin for IntelliJ IDEA.
 ## [Unreleased]
 
 ### Added
+* [INTELLIJ-194](https://jira.mongodb.org/browse/INTELLIJ-194) Add support for $addFields.
 * [INTELLIJ-193](https://jira.mongodb.org/browse/INTELLIJ-193) Add support for generating aggregates with $match and $project.
 * [INTELLIJ-189](https://jira.mongodb.org/browse/INTELLIJ-189) Add support for generating update queries.
 * [INTELLIJ-177](https://jira.mongodb.org/browse/INTELLIJ-177) Add support for parsing, inspecting and autocompleting in a addFields stage written using `Aggregation.addFields` and chained `AddFieldsOperation`s using `addFieldWithValue`, `addFieldWithValueOf`, `addField().withValue()` and `addField().withValueOf()`. Parsing boxed Java values is not supported yet.

diff --git a/...-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/AddFields.kt b/...-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/AddFields.kt
@@ -0,0 +1,19 @@
+package com.mongodb.jbplugin.dialects.mongosh.aggr
+
+import com.mongodb.jbplugin.dialects.mongosh.backend.MongoshBackend
+import com.mongodb.jbplugin.mql.Node
+import com.mongodb.jbplugin.mql.components.HasAddedFields
+
+internal fun <S> MongoshBackend.emitAddFieldsStage(node: Node<S>): MongoshBackend {
+    val addedFields = node.component<HasAddedFields<S>>()?.children ?: emptyList()
+    val isLongFieldList = addedFields.size > 3
+
+    emitObjectStart(long = isLongFieldList)
+    emitObjectKey(registerConstant('$' + "addFields"))
+    emitObjectStart(long = isLongFieldList)
+    emitAsFieldValueDocument(addedFields, isLongFieldList)
+    emitObjectEnd(long = isLongFieldList)
+    emitObjectEnd(long = isLongFieldList)
+
+    return this
+}
diff --git a/...-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Aggregate.kt b/...-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Aggregate.kt
@@ -33,30 +33,32 @@ fun <S> MongoshBackend.emitAggregateBody(node: Node<S>, queryContext: QueryConte
         when (stage.component<Named>()?.name) {
             Name.MATCH -> emitMatchStage(stage)
             Name.PROJECT -> emitProjectStage(stage)
+            Name.ADD_FIELDS -> emitAddFieldsStage(stage)
             else -> {}
         }
-        emitObjectValueEnd()
+        emitObjectValueEnd(long = true)
     }
     emitArrayEnd(long = true)
     return this
 }
 
-internal fun <S> MongoshBackend.emitAsFieldValueDocument(nodes: List<Node<S>>): MongoshBackend {
+internal fun <S> MongoshBackend.emitAsFieldValueDocument(nodes: List<Node<S>>, isLong: Boolean = false): MongoshBackend {
     for (node in nodes) {
         val field = node.component<HasFieldReference<S>>() ?: continue
         val value = node.component<HasValueReference<S>>() ?: continue
 
         emitObjectKey(resolveFieldReference(field))
         emitContextValue(resolveValueReference(value, field))
-        emitObjectValueEnd()
+        emitObjectValueEnd(long = isLong)
     }
 
     return this
 }
 
 private val NON_DESTRUCTIVE_STAGES = setOf(
     Name.MATCH,
-    Name.PROJECT
+    Name.PROJECT,
+    Name.ADD_FIELDS
 )
 
 private fun <S> Node<S>.isNotDestructive(): Boolean {

diff --git a/...db-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Project.kt b/...db-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Project.kt
@@ -11,7 +11,7 @@ internal fun <S> MongoshBackend.emitProjectStage(node: Node<S>): MongoshBackend
     emitObjectStart(long = isLongProjection)
     emitObjectKey(registerConstant('$' + "project"))
     emitObjectStart(long = isLongProjection)
-    emitAsFieldValueDocument(projections)
+    emitAsFieldValueDocument(projections, isLongProjection)
     emitObjectEnd(long = isLongProjection)
     emitObjectEnd(long = isLongProjection)
 

diff --git a/...s/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/backend/MongoshBackend.kt b/...s/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/backend/MongoshBackend.kt
@@ -114,8 +114,11 @@ class MongoshBackend(
         return this
     }
 
-    fun emitObjectValueEnd(): MongoshBackend {
+    fun emitObjectValueEnd(long: Boolean = false): MongoshBackend {
         emitAsIs(", ")
+        if (long && prettyPrint) {
+            emitNewLine()
+        }
         return this
     }
 

diff --git a/...lects/mongosh/src/test/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/AddFieldsTest.kt b/...lects/mongosh/src/test/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/AddFieldsTest.kt
@@ -0,0 +1,117 @@
+package com.mongodb.jbplugin.dialects.mongosh.aggr
+
+import com.mongodb.jbplugin.dialects.mongosh.assertGeneratedQuery
+import com.mongodb.jbplugin.mql.BsonInt32
+import com.mongodb.jbplugin.mql.Node
+import com.mongodb.jbplugin.mql.components.HasAddedFields
+import com.mongodb.jbplugin.mql.components.HasAggregation
+import com.mongodb.jbplugin.mql.components.HasFieldReference
+import com.mongodb.jbplugin.mql.components.HasValueReference
+import com.mongodb.jbplugin.mql.components.IsCommand
+import com.mongodb.jbplugin.mql.components.Name
+import com.mongodb.jbplugin.mql.components.Named
+import org.junit.jupiter.api.Test
+
+class AddFieldsTest {
+    @Test
+    fun `can format an addFields stage including fields`() {
+        assertGeneratedQuery(
+            """
+            var collection = ""
+            var database = ""
+
+            db.getSiblingDB(database).getCollection(collection).aggregate([{"${"$"}addFields": {"myField": 1, }}, ])
+            """.trimIndent()
+        ) {
+            Node(
+                Unit,
+                listOf(
+                    IsCommand(IsCommand.CommandType.AGGREGATE),
+                    HasAggregation(
+                        listOf(
+                            Node(
+                                Unit,
+                                listOf(
+                                    Named(Name.ADD_FIELDS),
+                                    HasAddedFields(
+                                        listOf(
+                                            Node(
+                                                Unit,
+                                                listOf(
+                                                    HasFieldReference(
+                                                        HasFieldReference.FromSchema(
+                                                            Unit,
+                                                            "myField"
+                                                        )
+                                                    ),
+                                                    HasValueReference(
+                                                        HasValueReference.Inferred(
+                                                            Unit,
+                                                            1,
+                                                            BsonInt32
+                                                        )
+                                                    )
+                                                )
+                                            )
+                                        )
+                                    )
+                                )
+                            )
+                        )
+                    )
+                )
+            )
+        }
+    }
+
+    @Test
+    fun `can format an addFields stage excluding fields`() {
+        assertGeneratedQuery(
+            """
+            var collection = ""
+            var database = ""
+
+            db.getSiblingDB(database).getCollection(collection).aggregate([{"${"$"}addFields": {"myField": -1, }}, ])
+            """.trimIndent()
+        ) {
+            Node(
+                Unit,
+                listOf(
+                    IsCommand(IsCommand.CommandType.AGGREGATE),
+                    HasAggregation(
+                        listOf(
+                            Node(
+                                Unit,
+                                listOf(
+                                    Named(Name.ADD_FIELDS),
+                                    HasAddedFields(
+                                        listOf(
+                                            Node(
+                                                Unit,
+                                                listOf(
+                                                    HasFieldReference(
+                                                        HasFieldReference.FromSchema(
+                                                            Unit,
+                                                            "myField"
+                                                        )
+                                                    ),
+                                                    HasValueReference(
+                                                        HasValueReference.Inferred(
+                                                            Unit,
+                                                            -1,
+                                                            BsonInt32
+                                                        )
+                                                    )
+                                                )
+                                            )
+                                        )
+                                    )
+                                )
+                            )
+                        )
+                    )
+                )
+            )
+        }
+    }
+}