feat(query-generation): support for the $project stage INTELLIJ-193 (#…

…126)
mongodb · Jan 23, 2025 · 99f18ef · 99f18ef
1 parent 8857817
commit 99f18ef
Show file tree

Hide file tree

Showing 12 changed files with 548 additions and 251 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@ MongoDB plugin for IntelliJ IDEA.
 ## [Unreleased]
 
 ### Added
+* [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.
 * [INTELLIJ-174](https://jira.mongodb.org/browse/INTELLIJ-174) Add support for parsing, inspecting and autocompleting in a sort stage written using `Aggregation.sort` and chained `SortOperation`s using `and`. All the overloads of creating a `Sort` object are supported.

diff --git a/...-plugin/src/main/kotlin/com/mongodb/jbplugin/codeActions/impl/RunQueryCodeActionBridge.kt b/...-plugin/src/main/kotlin/com/mongodb/jbplugin/codeActions/impl/RunQueryCodeActionBridge.kt
@@ -37,12 +37,6 @@ import com.mongodb.jbplugin.meta.service
 import com.mongodb.jbplugin.mql.Node
 import com.mongodb.jbplugin.mql.QueryContext
 import com.mongodb.jbplugin.mql.components.HasSourceDialect
-import com.mongodb.jbplugin.mql.components.IsCommand
-import com.mongodb.jbplugin.mql.parser.components.whenHasAnyCommand
-import com.mongodb.jbplugin.mql.parser.components.whenIsCommand
-import com.mongodb.jbplugin.mql.parser.first
-import com.mongodb.jbplugin.mql.parser.map
-import com.mongodb.jbplugin.mql.parser.parse
 import com.mongodb.jbplugin.observability.TelemetryEvent
 import com.mongodb.jbplugin.observability.TelemetryEvent.QueryRunEvent.Console
 import com.mongodb.jbplugin.observability.probe.QueryRunProbe
@@ -69,10 +63,6 @@ internal object RunQueryCodeAction : MongoDbCodeAction {
         query: Node<PsiElement>,
         formatter: DialectFormatter
     ): LineMarkerInfo<PsiElement>? {
-        if (!shouldShowRunGutterIcon(query)) {
-            return null
-        }
-
         return LineMarkerInfo(
             query.sourceForMarker,
             query.sourceForMarker.textRange,
@@ -124,14 +114,6 @@ internal object RunQueryCodeAction : MongoDbCodeAction {
         )
     }
 
-    private fun shouldShowRunGutterIcon(node: Node<PsiElement>): Boolean {
-        return first(
-            whenIsCommand<PsiElement>(IsCommand.CommandType.AGGREGATE).map { false },
-            whenHasAnyCommand<PsiElement>().map { true }
-        ).parse(node) // if we couldn't parse the query, don't show the gutter icon
-            .orElse { false }
-    }
-
     private fun openDataGripConsole(
         query: Node<PsiElement>,
         newDataSource: LocalDataSource,

diff --git a/...src/test/kotlin/com/mongodb/jbplugin/codeActions/impl/JavaDriverRunQueryCodeActionTest.kt b/...src/test/kotlin/com/mongodb/jbplugin/codeActions/impl/JavaDriverRunQueryCodeActionTest.kt
@@ -84,39 +84,4 @@ public class Repository {
         assertEquals(Icons.runQueryGutter, gutter.icon)
         assertEquals(CodeActionsMessages.message("code.action.run.query"), gutter.tooltipText)
     }
-
-    @ParsingTest(
-        fileName = "Repository.java",
-        value = """
-import com.mongodb.client.FindIterable;
-import com.mongodb.client.MongoClient;
-import com.mongodb.client.MongoCollection;
-import org.bson.Document;
-import org.bson.types.ObjectId;
-import java.util.List;import static com.mongodb.client.model.Filters.*;
-
-public class Repository {
-    private final MongoClient client;
-
-    public Repository(MongoClient client) {
-        this.client = client;
-    }
-
-    public FindIterable<Document> exampleFind() {
-        return client.getDatabase("myDatabase")
-                .getCollection("myCollection")
-                .aggregate(List.of()).first();
-    }
-}
-        """,
-    )
-    fun `does not show a gutter icon for aggregates`(
-        fixture: CodeInsightTestFixture,
-    ) {
-        fixture.setupConnection()
-        fixture.specifyDialect(JavaDriverDialect)
-
-        val gutters = fixture.findAllGutters()
-        assertEquals(0, gutters.size)
-    }
 }
diff --git a/...-java-project-with-mongodb/src/main/java/alt/mongodb/javadriver/JavaDriverRepository.java b/...-java-project-with-mongodb/src/main/java/alt/mongodb/javadriver/JavaDriverRepository.java
@@ -8,6 +8,7 @@
 import com.mongodb.client.model.Projections;
 import com.mongodb.client.model.Sorts;
 import org.bson.Document;
+import org.bson.types.ObjectId;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -57,13 +58,13 @@ public List<Document> findMoviesByYear(int year) {
             .into(new ArrayList<>());
     }
 
-    public Document queryMovieById(String id) {
+    public Document queryMovieById(ObjectId id) {
         return client
             .getDatabase("sample_mflix")
             .getCollection("movies")
             .aggregate(List.of(Aggregates.match(
                 Filters.eq(id)
-            ), Aggregates.unwind("year")))
+            ), Aggregates.project(Projections.include("title", "year"))))
             .first();
     }
 

diff --git a/.../mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/MongoshDialectFormatter.kt b/.../mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/MongoshDialectFormatter.kt
@@ -2,7 +2,6 @@ package com.mongodb.jbplugin.dialects.mongosh
 
 import com.mongodb.jbplugin.dialects.DialectFormatter
 import com.mongodb.jbplugin.dialects.OutputQuery
-import com.mongodb.jbplugin.dialects.mongosh.aggr.canEmitAggregate
 import com.mongodb.jbplugin.dialects.mongosh.aggr.emitAggregateBody
 import com.mongodb.jbplugin.dialects.mongosh.aggr.isAggregate
 import com.mongodb.jbplugin.dialects.mongosh.backend.MongoshBackend
@@ -23,14 +22,8 @@ object MongoshDialectFormatter : DialectFormatter {
         queryContext: QueryContext,
     ): OutputQuery {
         val isAggregate = query.isAggregate()
-        val canEmitAggregate = query.canEmitAggregate()
 
         val outputString = MongoshBackend(prettyPrint = queryContext.prettyPrint).apply {
-            if (isAggregate && !canEmitAggregate) {
-                emitComment("Only aggregates with a single match stage can be converted.")
-                return@apply
-            }
-
             emitDbAccess()
             emitCollectionReference(query.component<HasCollectionReference<S>>())
             if (queryContext.explainPlan != QueryContext.ExplainPlanType.NONE) {
@@ -62,7 +55,7 @@ object MongoshDialectFormatter : DialectFormatter {
             } else {
                 emitFunctionCall(long = true, {
                     if (query.isAggregate()) {
-                        emitAggregateBody(query)
+                        emitAggregateBody(query, queryContext)
                     } else {
                         emitQueryFilter(query, firstCall = true)
                     }
@@ -76,7 +69,6 @@ object MongoshDialectFormatter : DialectFormatter {
 
         val ref = query.component<HasCollectionReference<S>>()?.reference
         return when {
-            isAggregate && !canEmitAggregate -> OutputQuery.Incomplete(outputString)
             ref is HasCollectionReference.Known -> if (ref.namespace.isValid) {
                 OutputQuery.CanBeRun(outputString)
             } else {

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
@@ -1,21 +1,18 @@
 package com.mongodb.jbplugin.dialects.mongosh.aggr
 
 import com.mongodb.jbplugin.dialects.mongosh.backend.MongoshBackend
-import com.mongodb.jbplugin.dialects.mongosh.query.emitQueryFilter
+import com.mongodb.jbplugin.dialects.mongosh.query.resolveFieldReference
+import com.mongodb.jbplugin.dialects.mongosh.query.resolveValueReference
 import com.mongodb.jbplugin.mql.Node
+import com.mongodb.jbplugin.mql.QueryContext
 import com.mongodb.jbplugin.mql.components.HasAggregation
-import com.mongodb.jbplugin.mql.components.HasFilter
+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.parser.anyError
-import com.mongodb.jbplugin.mql.parser.components.aggregationStages
-import com.mongodb.jbplugin.mql.parser.components.hasName
+import com.mongodb.jbplugin.mql.components.Named
 import com.mongodb.jbplugin.mql.parser.components.whenIsCommand
-import com.mongodb.jbplugin.mql.parser.count
-import com.mongodb.jbplugin.mql.parser.filter
 import com.mongodb.jbplugin.mql.parser.map
-import com.mongodb.jbplugin.mql.parser.matches
-import com.mongodb.jbplugin.mql.parser.nth
 import com.mongodb.jbplugin.mql.parser.parse
 
 fun <S> Node<S>.isAggregate(): Boolean {
@@ -24,33 +21,44 @@ fun <S> Node<S>.isAggregate(): Boolean {
         .parse(this).orElse { false }
 }
 
-fun <S> Node<S>.canEmitAggregate(): Boolean {
-    return aggregationStages<S>()
-        .matches(count<Node<S>>().filter { it >= 1 }.matches().anyError())
-        .nth(0)
-        .matches(hasName(Name.MATCH))
-        .map { true }
-        .parse(this).orElse { false }
-}
-
-fun <S> MongoshBackend.emitAggregateBody(node: Node<S>): MongoshBackend {
-    // here we can assume that we only have 1 single stage that is a match
-    val matchStage = node.component<HasAggregation<S>>()!!.children[0]
-    val filter = matchStage.component<HasFilter<S>>()?.children?.getOrNull(0)
-    val longFilter = filter?.component<HasFilter<S>>()?.children?.size?.let { it > 3 } == true
+fun <S> MongoshBackend.emitAggregateBody(node: Node<S>, queryContext: QueryContext): MongoshBackend {
+    val allStages = node.component<HasAggregation<S>>()?.children ?: emptyList()
+    val stagesToEmit = when (queryContext.explainPlan) {
+        QueryContext.ExplainPlanType.NONE -> allStages
+        else -> allStages.takeWhile { it.isNotDestructive() }
+    }
 
     emitArrayStart(long = true)
-    emitObjectStart()
-    emitObjectKey(registerConstant('$' + "match"))
-    if (filter != null) {
-        emitObjectStart(long = longFilter)
-        emitQueryFilter(filter)
-        emitObjectEnd(long = longFilter)
-    } else {
-        emitComment("No filter provided.")
+    for (stage in stagesToEmit) {
+        when (stage.component<Named>()?.name) {
+            Name.MATCH -> emitMatchStage(stage)
+            Name.PROJECT -> emitProjectStage(stage)
+            else -> {}
+        }
+        emitObjectValueEnd()
     }
-    emitObjectEnd()
     emitArrayEnd(long = true)
+    return this
+}
+
+internal fun <S> MongoshBackend.emitAsFieldValueDocument(nodes: List<Node<S>>): 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()
+    }
 
     return this
 }
+
+private val NON_DESTRUCTIVE_STAGES = setOf(
+    Name.MATCH,
+    Name.PROJECT
+)
+
+private fun <S> Node<S>.isNotDestructive(): Boolean {
+    return component<Named>()?.name?.let { NON_DESTRUCTIVE_STAGES.contains(it) } == true
+}
diff --git a/...godb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Match.kt b/...godb-dialects/mongosh/src/main/kotlin/com/mongodb/jbplugin/dialects/mongosh/aggr/Match.kt
@@ -0,0 +1,24 @@
+package com.mongodb.jbplugin.dialects.mongosh.aggr
+
+import com.mongodb.jbplugin.dialects.mongosh.backend.MongoshBackend
+import com.mongodb.jbplugin.dialects.mongosh.query.emitQueryFilter
+import com.mongodb.jbplugin.mql.Node
+import com.mongodb.jbplugin.mql.components.HasFilter
+
+internal fun <S> MongoshBackend.emitMatchStage(node: Node<S>): MongoshBackend {
+    val filter = node.component<HasFilter<S>>()?.children?.getOrNull(0)
+    val longFilter = filter?.component<HasFilter<S>>()?.children?.size?.let { it > 3 } == true
+
+    emitObjectStart()
+    emitObjectKey(registerConstant('$' + "match"))
+    if (filter != null) {
+        emitObjectStart(long = longFilter)
+        emitQueryFilter(filter)
+        emitObjectEnd(long = longFilter)
+    } else {
+        emitComment("No filter provided.")
+    }
+    emitObjectEnd()
+
+    return this
+}
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
@@ -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.HasProjections
+
+internal fun <S> MongoshBackend.emitProjectStage(node: Node<S>): MongoshBackend {
+    val projections = node.component<HasProjections<S>>()?.children ?: emptyList()
+    val isLongProjection = projections.size > 3
+
+    emitObjectStart(long = isLongProjection)
+    emitObjectKey(registerConstant('$' + "project"))
+    emitObjectStart(long = isLongProjection)
+    emitAsFieldValueDocument(projections)
+    emitObjectEnd(long = isLongProjection)
+    emitObjectEnd(long = isLongProjection)
+
+    return this
+}