feat(query-generation): support for generating update queries INTELLIJ-189

CHANGELOG.md

@@ -5,6 +5,7 @@ MongoDB plugin for IntelliJ IDEA.
 ## [Unreleased]
 
 ### Added
+* [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.
 * [INTELLIJ-188](https://jira.mongodb.org/browse/INTELLIJ-188) Support for generating sort in the query generator.

packages/jetbrains-plugin/src/main/kotlin/com/mongodb/jbplugin/codeActions/impl/RunQueryCodeActionBridge.kt

@@ -87,7 +87,7 @@ internal object RunQueryCodeAction : MongoDbCodeAction {
                     coroutineScope.launchChildBackground {
                         val outputQuery = MongoshDialect.formatter.formatQuery(
                             query,
-                            QueryContext.empty()
+                            QueryContext.empty(prettyPrint = true)
                         )
 
                         if (dataSource?.isConnected() == true) {

packages/jetbrains-plugin/src/main/kotlin/com/mongodb/jbplugin/inspections/impl/IndexCheckInspectionBridge.kt

@@ -66,7 +66,8 @@ internal object IndexCheckLinterInspection : MongoDbInspection {
                 QueryContext.ExplainPlanType.FULL
             } else {
                 QueryContext.ExplainPlanType.SAFE
-            }
+            },
+            prettyPrint = false
         )
 
         val readModelProvider by query.source.project.service<DataGripBasedReadModelProvider>()

packages/jetbrains-plugin/src/test/resources/project-fixtures/basic-java-project-with-mongodb/src/main/java/alt/mongodb/javadriver/JavaDriverRepository.java

@@ -30,6 +30,23 @@ public Document findMovieById(String id) {
             .first();
     }
 
+    public void updateMoviesByYear(int year) {
+        client
+            .getDatabase("sample_mflix")
+            .getCollection("movies")
+            .updateMany(
+                Filters.and(Filters.eq("year", year), Filters.ne("languages", "Esperanto")),
+                Updates.combine(
+                    Updates.inc(IMDB_VOTES, 1),
+                    Updates.inc(AWARDS_WINS, 1),
+                    Updates.set("other", 25),
+                    Updates.pull(IMDB_VOTES, Filters.eq("a", 1)),
+                    Updates.pullAll(IMDB_VOTES, List.of("1", "2")),
+                    Updates.push("languages", "Esperanto")
+                )
+            );
+    }
+
     public List<Document> findMoviesByYear(int year) {
         return client
             .getDatabase("sample_mflix")

packages/mongodb-access-adapter/datagrip-access-adapter/src/test/kotlin/com/mongodb/jbplugin/accessadapter/datagrip/adapter/DataGripMongoDbDriverTest.kt

@@ -273,7 +273,7 @@ class DataGripMongoDbDriverTest {
 
         val explainPlanResult = driver.explain(
             query,
-            QueryContext(emptyMap(), QueryContext.ExplainPlanType.SAFE)
+            QueryContext(emptyMap(), QueryContext.ExplainPlanType.SAFE, false)
         )
         assertEquals(ExplainPlan.CollectionScan, explainPlanResult)
     }
@@ -324,7 +324,7 @@ class DataGripMongoDbDriverTest {
 
         val explainPlanResult = driver.explain(
             query,
-            QueryContext(emptyMap(), QueryContext.ExplainPlanType.SAFE)
+            QueryContext(emptyMap(), QueryContext.ExplainPlanType.SAFE, false)
         )
         assertEquals(ExplainPlan.IndexScan, explainPlanResult)
     }

packages/mongodb-dialects/java-driver/src/main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParser.kt

@@ -392,7 +392,7 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
                 filter,
                 listOf(
                     Named(Name.from(method.name)),
-                    HasFilter(
+                    HasUpdates(
                         filter.argumentList.expressions
                             .mapNotNull { resolveBsonBuilderCall(it, UPDATES_FQN) }
                             .mapNotNull { parseUpdatesExpression(it) },
@@ -402,12 +402,36 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
         } else if (method.parameters.size == 2) {
             // If it has two parameters, it's field/value.
             val fieldReference = filter.argumentList.expressions[0].resolveFieldNameFromExpression()
+            val named = Named(Name.from(method.name))
+
+            if (named.name == Name.PULL) {
+                // the second argument can be either a value or a filter
+                val filterExpression = filter.argumentList.expressions[1]
+                val resolvedFilterExpression = resolveBsonBuilderCall(filterExpression, FILTERS_FQN)
+
+                if (resolvedFilterExpression != null) {
+                    val parsedFilter = parseFilterExpression(resolvedFilterExpression)
+                        ?: return null
+
+                    return Node(
+                        filter,
+                        listOf(
+                            named,
+                            HasFieldReference(
+                                fieldReference,
+                            ),
+                            HasFilter(listOf(parsedFilter)),
+                        ),
+                    )
+                }
+            }
+
             val valueReference = resolveValueFromExpression(filter.argumentList.expressions[1])
 
             return Node(
                 filter,
                 listOf(
-                    Named(Name.from(method.name)),
+                    named,
                     HasFieldReference(
                         fieldReference,
                     ),

packages/mongodb-dialects/java-driver/src/main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/PsiMdbTreeUtil.kt

@@ -6,13 +6,41 @@ package com.mongodb.jbplugin.dialects.javadriver.glossary
 
 import com.intellij.lang.jvm.JvmModifier
 import com.intellij.openapi.project.Project
-import com.intellij.psi.*
+import com.intellij.psi.JavaPsiFacade
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiExpression
+import com.intellij.psi.PsiField
+import com.intellij.psi.PsiLiteralExpression
+import com.intellij.psi.PsiLiteralValue
+import com.intellij.psi.PsiLocalVariable
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiMethodCallExpression
+import com.intellij.psi.PsiParenthesizedExpression
+import com.intellij.psi.PsiReference
+import com.intellij.psi.PsiReferenceExpression
+import com.intellij.psi.PsiReturnStatement
+import com.intellij.psi.PsiSuperExpression
+import com.intellij.psi.PsiThisExpression
+import com.intellij.psi.PsiType
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.PsiTypesUtil
 import com.intellij.psi.util.childrenOfType
 import com.intellij.psi.util.parentOfType
-import com.mongodb.jbplugin.mql.*
+import com.mongodb.jbplugin.mql.BsonAny
+import com.mongodb.jbplugin.mql.BsonAnyOf
+import com.mongodb.jbplugin.mql.BsonArray
+import com.mongodb.jbplugin.mql.BsonBoolean
+import com.mongodb.jbplugin.mql.BsonDate
+import com.mongodb.jbplugin.mql.BsonDecimal128
+import com.mongodb.jbplugin.mql.BsonDouble
+import com.mongodb.jbplugin.mql.BsonInt32
+import com.mongodb.jbplugin.mql.BsonInt64
+import com.mongodb.jbplugin.mql.BsonNull
+import com.mongodb.jbplugin.mql.BsonObjectId
+import com.mongodb.jbplugin.mql.BsonString
+import com.mongodb.jbplugin.mql.BsonType
 import com.mongodb.jbplugin.mql.components.IsCommand
 
 /**
@@ -398,6 +426,9 @@ fun String.toBsonType(): BsonType {
     } else if (this.endsWith("[]")) {
         val baseType = this.substring(0, this.length - 2)
         return BsonArray(baseType.toBsonType())
+    } else if (this.contains("List") || this.contains("Set")) {
+        val baseType = this.substringAfter("<").substringBeforeLast(">")
+        return BsonArray(baseType.toBsonType())
     }
 
     return BsonAny

packages/mongodb-dialects/java-driver/src/test/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParserTest.kt

@@ -1709,9 +1709,9 @@ public class Repository {
 
         val combine = hasUpdate.children[0]
         assertEquals(Name.COMBINE, combine.component<Named>()!!.name)
-        assertEquals(2, combine.component<HasFilter<Unit?>>()!!.children.size)
+        assertEquals(2, combine.component<HasUpdates<Unit?>>()!!.children.size)
 
-        val updates = combine.component<HasFilter<Unit?>>()!!.children
+        val updates = combine.component<HasUpdates<Unit?>>()!!.children
         assertEquals(Name.SET, updates[0].component<Named>()!!.name)
         assertEquals(Name.UNSET, updates[1].component<Named>()!!.name)
     }

packages/mongodb-dialects/java-driver/src/test/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/PsiMdbTreeUtilTest.kt

@@ -28,6 +28,15 @@ class PsiMdbTreeUtilTest {
         }
     }
 
+    @ParameterizedTest
+    @MethodSource("stringToBsonType")
+    fun `should map all known java qualified names to their corresponding bson types`(
+        javaQualifiedName: String,
+        expected: BsonType,
+    ) {
+        assertEquals(expected, javaQualifiedName.toBsonType())
+    }
+
     companion object {
         @JvmStatic
         fun psiTypeToBsonType(): Array<Array<Any>> =
@@ -89,6 +98,38 @@ class PsiMdbTreeUtilTest {
                     BsonAnyOf(BsonDecimal128, BsonNull),
                 ),
             )
+
+        @JvmStatic
+        fun stringToBsonType(): Array<Array<Any>> = arrayOf(
+            arrayOf("org.bson.types.ObjectId", BsonAnyOf(BsonObjectId, BsonNull)),
+            arrayOf("boolean", BsonBoolean),
+            arrayOf("java.lang.Boolean", BsonBoolean),
+            arrayOf("short", BsonInt32),
+            arrayOf("java.lang.Short", BsonInt32),
+            arrayOf("int", BsonInt32),
+            arrayOf("java.lang.Integer", BsonInt32),
+            arrayOf("long", BsonInt64),
+            arrayOf("java.lang.Long", BsonInt64),
+            arrayOf("float", BsonDouble),
+            arrayOf("java.lang.Float", BsonDouble),
+            arrayOf("double", BsonDouble),
+            arrayOf("java.lang.Double", BsonDouble),
+            arrayOf("java.lang.CharSequence", BsonAnyOf(BsonString, BsonNull)),
+            arrayOf("java.lang.String", BsonAnyOf(BsonString, BsonNull)),
+            arrayOf("String", BsonAnyOf(BsonString, BsonNull)),
+            arrayOf("java.util.Date", BsonAnyOf(BsonDate, BsonNull)),
+            arrayOf("java.time.LocalDate", BsonAnyOf(BsonDate, BsonNull)),
+            arrayOf("java.time.LocalDateTime", BsonAnyOf(BsonDate, BsonNull)),
+            arrayOf("java.math.BigInteger", BsonAnyOf(BsonInt64, BsonNull)),
+            arrayOf("java.math.BigDecimal", BsonAnyOf(BsonDecimal128, BsonNull)),
+            arrayOf("int[]", BsonArray(BsonInt32)),
+            arrayOf("java.lang.Long[]", BsonArray(BsonInt64)),
+            arrayOf("List<String>", BsonArray(BsonAnyOf(BsonString, BsonNull))),
+            arrayOf("List<java.lang.Integer>[]", BsonArray(BsonArray(BsonInt32))),
+            arrayOf("Set<String>", BsonArray(BsonAnyOf(BsonString, BsonNull))),
+            arrayOf("Map<String, Integer>", BsonAny),
+            arrayOf("HashMap<String, Integer>", BsonAny),
+        )
     }
 }