chore: fix issue with aggregate.first

mongodb · Nov 14, 2024 · d4afb0f · d4afb0f
1 parent f9e65bd
commit d4afb0f
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 4 deletions.
diff --git a/.../main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParser.kt b/.../main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParser.kt
@@ -544,7 +544,7 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
             )
             .flatMap(methodCallChain()).acceptAnyError()
             .firstMatching(
-                resolveFromMethodName.flatMap(equalsTo(IsCommand.CommandType.AGGREGATE)).matches()
+                resolveFromMethodName.filter { it == IsCommand.CommandType.AGGREGATE }.matches()
             ).flatMap(resolveFromMethodName)
             .acceptAnyError()
 

diff --git a/...river/src/main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/parser/methods.kt b/...river/src/main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/parser/methods.kt
@@ -6,6 +6,7 @@ import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiMethodCallExpression
 import com.intellij.psi.PsiReturnStatement
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.util.parentsOfType
 import com.mongodb.jbplugin.dialects.javadriver.glossary.collectTypeUntil
 import com.mongodb.jbplugin.dialects.javadriver.glossary.findAllChildrenOfType
 import com.mongodb.jbplugin.dialects.javadriver.glossary.fuzzyResolveMethod
@@ -85,7 +86,8 @@ fun methodCallChain(): Parser<PsiMethodCallExpression, Any, List<PsiMethodCallEx
         // by reversing the list, we get the "deepest" (closest) levels first
         val allCallExpressions = input.findAllChildrenOfType(
             PsiMethodCallExpression::class.java
-        ).reversed()
+        ).reversed() + input.parentsOfType<PsiMethodCallExpression>()
+
         Either.right(allCallExpressions)
     }
 }

diff --git a/...t/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParserTest.kt b/...t/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParserTest.kt
@@ -136,7 +136,7 @@ public final class Repository {
         // Returns the value of the entire return expression
         val queryWithIterableCall = psiFile.getQueryAtMethod("Repository", "findBookById")
         // Now the candidate for query attaches at the top element
-        assertTrue(JavaDriverDialectParser.isCandidateForQuery(queryWithIterableCall))
+        assertFalse(JavaDriverDialectParser.isCandidateForQuery(queryWithIterableCall))
 
         val actualQuery = PsiTreeUtil
             .findChildrenOfType(queryWithIterableCall, PsiMethodCallExpression::class.java)
@@ -1652,4 +1652,39 @@ public final class Repository {
         val command = parsedQuery.component<IsCommand>()
         assertEquals(IsCommand.CommandType.FIND_ONE, command?.type)
     }
+
+    @ParsingTest(
+        fileName = "Repository.java",
+        value = """
+import com.mongodb.client.FindIterable;
+import com.mongodb.client.MongoClient;
+import org.bson.Document;
+import java.util.List;
+import static com.mongodb.client.model.Filters.*;
+import static com.mongodb.client.model.Updates.*;
+
+public class Repository {
+    private final MongoClient client;
+
+    public Repository(MongoClient client) {
+        this.client = client;
+    }
+    
+    private FindIterable<Document> findFantasyBooks() {
+        return client.getDatabase("myDatabase")
+                .getCollection("myCollection")
+                .aggregate(List.of()).first();
+    }
+}
+        """,
+    )
+    fun `understands aggregate first as an AGGREGATE command`(
+        psiFile: PsiFile
+    ) {
+        val query = psiFile.getQueryAtMethod("Repository", "findFantasyBooks")
+        val parsedQuery = JavaDriverDialect.parser.parse(query)
+
+        val command = parsedQuery.component<IsCommand>()
+        assertEquals(IsCommand.CommandType.AGGREGATE, command?.type)
+    }
 }
diff --git a/packages/mongodb-mql-model/src/main/kotlin/com/mongodb/jbplugin/mql/parser/Parser.kt b/packages/mongodb-mql-model/src/main/kotlin/com/mongodb/jbplugin/mql/parser/Parser.kt
@@ -178,7 +178,7 @@ inline fun <reified II, I, E, O> Parser<I, E, O>.inputAs(): Parser<II, E, O> {
 fun <I, E, O> Parser<I, E, O>.debug(message: String): Parser<I, E, O> {
     return { input ->
         val result = this(input)
-        println("$input > $result >> $message ")
+        println(":: $message -> $input > $result ")
         result
     }
 }