From 4433d1b9b5b556819f6ee29e9a1a66443f3f1e3a Mon Sep 17 00:00:00 2001 From: Kevin Mas Ruiz Date: Fri, 24 Jan 2025 10:15:05 +0100 Subject: [PATCH] feat(query-generation): add support for $sort INTELLIJ-196 (#131) --- CHANGELOG.md | 1 + .../dialects/mongosh/aggr/Aggregate.kt | 4 +- .../jbplugin/dialects/mongosh/aggr/Sort.kt | 19 +++++ .../dialects/mongosh/query/References.kt | 35 +++++--- .../dialects/mongosh/aggr/SortTest.kt | 84 +++++++++++++++++++ 5 files changed, 130 insertions(+), 13 deletions(-) create mode 100644 packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Sort.kt create mode 100644 packages/mongodb-dialects/mongosh/src/test/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/SortTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index fe1c5461..79e5c926 100644 --- 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/packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Aggregate.kt b/packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Aggregate.kt index f022c722..7c1bbcad 100644 --- a/packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Aggregate.kt +++ b/packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Aggregate.kt @@ -35,6 +35,7 @@ fun MongoshBackend.emitAggregateBody(node: Node, 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 Node.isNotDestructive(): Boolean { diff --git a/packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Sort.kt b/packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Sort.kt new file mode 100644 index 00000000..dcf50500 --- /dev/null +++ b/packages/mongodb-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 MongoshBackend.emitSortStage(node: Node): MongoshBackend { + val sorts = node.component>()?.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/packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/query/References.kt b/packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/query/References.kt index 371777f3..405cb453 100644 --- a/packages/mongodb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/query/References.kt +++ b/packages/mongodb-dialects/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 MongoshBackend.resolveValueReference( valueRef: HasValueReference, fieldRef: HasFieldReference? -) = 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>() + ?: 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 MongoshBackend.resolveFieldReference(fieldRef: HasFieldReference) = diff --git a/packages/mongodb-dialects/mongosh/src/test/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/SortTest.kt b/packages/mongodb-dialects/mongosh/src/test/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/SortTest.kt new file mode 100644 index 00000000..baa48c59 --- /dev/null +++ b/packages/mongodb-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 + ) + ) + ) + ) + ) + ) + ) + ) + ) + ) + ) + ) + } + } +}