feat(query-generation): support for the unwind stages INTELLIJ-195 (#130

)
mongodb · Jan 23, 2025 · b0c0432 · b0c0432
1 parent d7678ff
commit b0c0432
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,7 +5,8 @@ MongoDB plugin for IntelliJ IDEA.
 ## [Unreleased]
 
 ### Added
-* [INTELLIJ-194](https://jira.mongodb.org/browse/INTELLIJ-194) Add support for $addFields.
+* [INTELLIJ-195](https://jira.mongodb.org/browse/INTELLIJ-195) Add support for $unwind when generating the query into DataGrip.
+* [INTELLIJ-194](https://jira.mongodb.org/browse/INTELLIJ-194) Add support for $addFields when generating the query into DataGrip.
 * [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/Aggregate.kt b/...-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Aggregate.kt
@@ -34,6 +34,7 @@ fun <S> MongoshBackend.emitAggregateBody(node: Node<S>, queryContext: QueryConte
             Name.MATCH -> emitMatchStage(stage)
             Name.PROJECT -> emitProjectStage(stage)
             Name.ADD_FIELDS -> emitAddFieldsStage(stage)
+            Name.UNWIND -> emitUnwindStage(stage)
             else -> {}
         }
         emitObjectValueEnd(long = true)
@@ -58,7 +59,8 @@ internal fun <S> MongoshBackend.emitAsFieldValueDocument(nodes: List<Node<S>>, i
 private val NON_DESTRUCTIVE_STAGES = setOf(
     Name.MATCH,
     Name.PROJECT,
-    Name.ADD_FIELDS
+    Name.ADD_FIELDS,
+    Name.UNWIND
 )
 
 private fun <S> Node<S>.isNotDestructive(): Boolean {

diff --git a/...odb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Unwind.kt b/...odb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Unwind.kt
@@ -0,0 +1,17 @@
+package com.mongodb.jbplugin.dialects.mongosh.aggr
+
+import com.mongodb.jbplugin.dialects.mongosh.backend.MongoshBackend
+import com.mongodb.jbplugin.dialects.mongosh.query.resolveFieldReference
+import com.mongodb.jbplugin.mql.Node
+import com.mongodb.jbplugin.mql.components.HasFieldReference
+
+internal fun <S> MongoshBackend.emitUnwindStage(node: Node<S>): MongoshBackend {
+    val unwindField = node.component<HasFieldReference<S>>() ?: return this
+
+    emitObjectStart()
+    emitObjectKey(registerConstant('$' + "unwind"))
+    emitContextValue(resolveFieldReference(unwindField))
+    emitObjectEnd()
+
+    return this
+}
diff --git a/...ialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/query/References.kt b/...ialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/query/References.kt
@@ -29,9 +29,9 @@ fun <S> MongoshBackend.resolveValueReference(
 
 fun <S> MongoshBackend.resolveFieldReference(fieldRef: HasFieldReference<S>) =
     when (val ref = fieldRef.reference) {
-        is FromSchema -> registerConstant(ref.fieldName)
-        is Inferred<S> -> registerConstant(ref.fieldName)
-        is HasFieldReference.Computed<S> -> registerConstant(ref.fieldName)
+        is FromSchema -> registerConstant(ref.displayName)
+        is Inferred<S> -> registerConstant(ref.displayName)
+        is HasFieldReference.Computed<S> -> registerConstant(ref.displayName)
         is Unknown -> registerVariable("field", BsonAny)
     }
 

diff --git a/...dialects/mongosh/src/test/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/UnwindTest.kt b/...dialects/mongosh/src/test/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/UnwindTest.kt
@@ -0,0 +1,47 @@
+package com.mongodb.jbplugin.dialects.mongosh.aggr
+
+import com.mongodb.jbplugin.dialects.mongosh.assertGeneratedQuery
+import com.mongodb.jbplugin.mql.Node
+import com.mongodb.jbplugin.mql.components.HasAggregation
+import com.mongodb.jbplugin.mql.components.HasFieldReference
+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 UnwindTest {
+    @Test
+    fun `can format an unwind stage`() {
+        assertGeneratedQuery(
+            """
+            var collection = ""
+            var database = ""
+
+            db.getSiblingDB(database).getCollection(collection).aggregate([{"${"$"}unwind": "${'$'}myField"}, ])
+            """.trimIndent()
+        ) {
+            Node(
+                Unit,
+                listOf(
+                    IsCommand(IsCommand.CommandType.AGGREGATE),
+                    HasAggregation(
+                        listOf(
+                            Node(
+                                Unit,
+                                listOf(
+                                    Named(Name.UNWIND),
+                                    HasFieldReference(
+                                        HasFieldReference.FromSchema(
+                                            Unit,
+                                            "${'$'}myField"
+                                        )
+                                    )
+                                )
+                            )
+                        )
+                    )
+                )
+            )
+        }
+    }
+}