feat(query-generation): add support for $sort INTELLIJ-196 (#131)

mongodb · Jan 24, 2025 · 4433d1b · 4433d1b
1 parent b0c0432
commit 4433d1b
Show file tree

Hide file tree

Showing 5 changed files with 130 additions and 13 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@ MongoDB plugin for IntelliJ IDEA.
 ## [Unreleased]
 
 ### Added
+* [INTELLIJ-196](https://jira.mongodb.org/browse/INTELLIJ-196) Add support for $sort when generating the query into DataGrip.
 * [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.

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
@@ -35,6 +35,7 @@ fun <S> MongoshBackend.emitAggregateBody(node: Node<S>, queryContext: QueryConte
             Name.PROJECT -> emitProjectStage(stage)
             Name.ADD_FIELDS -> emitAddFieldsStage(stage)
             Name.UNWIND -> emitUnwindStage(stage)
+            Name.SORT -> emitSortStage(stage)
             else -> {}
         }
         emitObjectValueEnd(long = true)
@@ -60,7 +61,8 @@ private val NON_DESTRUCTIVE_STAGES = setOf(
     Name.MATCH,
     Name.PROJECT,
     Name.ADD_FIELDS,
-    Name.UNWIND
+    Name.UNWIND,
+    Name.SORT,
 )
 
 private fun <S> Node<S>.isNotDestructive(): Boolean {

diff --git a/...ngodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Sort.kt b/...ngodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Sort.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.HasSorts
+
+internal fun <S> MongoshBackend.emitSortStage(node: Node<S>): MongoshBackend {
+    val sorts = node.component<HasSorts<S>>()?.children ?: emptyList()
+    val isLongSortChain = sorts.size > 3
+
+    emitObjectStart(long = isLongSortChain)
+    emitObjectKey(registerConstant('$' + "sort"))
+    emitObjectStart(long = isLongSortChain)
+    emitAsFieldValueDocument(sorts, isLongSortChain)
+    emitObjectEnd(long = isLongSortChain)
+    emitObjectEnd(long = isLongSortChain)
+
+    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
@@ -1,5 +1,6 @@
 package com.mongodb.jbplugin.dialects.mongosh.query
 
+import com.mongodb.jbplugin.dialects.mongosh.backend.ContextValue
 import com.mongodb.jbplugin.dialects.mongosh.backend.MongoshBackend
 import com.mongodb.jbplugin.mql.BsonAny
 import com.mongodb.jbplugin.mql.BsonString
@@ -13,18 +14,28 @@ import com.mongodb.jbplugin.mql.components.HasValueReference
 fun <S> MongoshBackend.resolveValueReference(
     valueRef: HasValueReference<S>,
     fieldRef: HasFieldReference<S>?
-) = when (val ref = valueRef.reference) {
-    is HasValueReference.Constant -> registerConstant(ref.value)
-    is HasValueReference.Inferred -> registerConstant(ref.value)
-    is HasValueReference.Runtime -> registerVariable(
-        (fieldRef?.reference as? FromSchema)?.fieldName ?: "value",
-        ref.type
-    )
-
-    else -> registerVariable(
-        "queryField",
-        BsonAny
-    )
+): ContextValue {
+    return when (val ref = valueRef.reference) {
+        is HasValueReference.Computed -> {
+            val node = ref.type.expression
+            val fieldRef =
+                node.component<HasFieldReference<S>>()
+                    ?: return registerVariable("queryField", ref.type.baseType)
+
+            resolveFieldReference(fieldRef)
+        }
+        is HasValueReference.Constant -> registerConstant(ref.value)
+        is HasValueReference.Inferred -> registerConstant(ref.value)
+        is HasValueReference.Runtime -> registerVariable(
+            (fieldRef?.reference as? FromSchema)?.fieldName ?: "value",
+            ref.type
+        )
+
+        else -> registerVariable(
+            "queryField",
+            BsonAny
+        )
+    }
 }
 
 fun <S> MongoshBackend.resolveFieldReference(fieldRef: HasFieldReference<S>) =

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