From 49b25781bcb1f1a6dc4fc8aa717a4f1b96b23881 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 16 Jan 2025 12:26:47 +0100 Subject: [PATCH 1/9] chore: refactor into dedicated aggregation parsers This PR splits the current aggregation parsing logic into dedicated stage based parsers. --- .../SpringCriteriaRepository.java | 7 +- .../glossary/JavaDriverDialectParser.kt | 32 ++++----- .../springcriteria/AggregationStagesParser.kt | 69 +++++++++++++++++++ .../SpringCriteriaDialectParser.kt | 53 ++------------ .../aggregationparsers/MatchStageParser.kt | 38 ++++++++++ ...Test.kt => AggregationStagesParserTest.kt} | 2 +- 6 files changed, 134 insertions(+), 67 deletions(-) create mode 100644 packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt create mode 100644 packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt rename packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/{AggregationParserTest.kt => AggregationStagesParserTest.kt} (99%) diff --git a/packages/jetbrains-plugin/src/test/resources/project-fixtures/basic-java-project-with-mongodb/src/main/java/alt/mongodb/springcriteria/SpringCriteriaRepository.java b/packages/jetbrains-plugin/src/test/resources/project-fixtures/basic-java-project-with-mongodb/src/main/java/alt/mongodb/springcriteria/SpringCriteriaRepository.java index 8d02cc7b..d3229d3b 100644 --- a/packages/jetbrains-plugin/src/test/resources/project-fixtures/basic-java-project-with-mongodb/src/main/java/alt/mongodb/springcriteria/SpringCriteriaRepository.java +++ b/packages/jetbrains-plugin/src/test/resources/project-fixtures/basic-java-project-with-mongodb/src/main/java/alt/mongodb/springcriteria/SpringCriteriaRepository.java @@ -2,7 +2,9 @@ import org.springframework.data.mongodb.core.MongoTemplate; import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields; import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Field; import java.util.List; @@ -32,9 +34,8 @@ private List allMoviesWithRatingAtLeast(int rating) { private List allMoviesWithRatingAtLeastAgg(int rating) { return template.aggregate( Aggregation.newAggregation( - Aggregation.match( - where( "tomatoes.viewer.rating").gte(rating) - ) + Aggregation.match(where( "tomatoes.viewer.rating").gte(rating)), + Aggregation.project("asd").andInclude("asd").andExclude("qwe") ), Movie.class, Movie.class diff --git a/packages/mongodb-dialects/java-driver/src/main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParser.kt b/packages/mongodb-dialects/java-driver/src/main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParser.kt index d60f294c..603ee161 100644 --- a/packages/mongodb-dialects/java-driver/src/main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParser.kt +++ b/packages/mongodb-dialects/java-driver/src/main/kotlin/com/mongodb/jbplugin/dialects/javadriver/glossary/JavaDriverDialectParser.kt @@ -267,7 +267,7 @@ object JavaDriverDialectParser : DialectParser { return null // empty, do nothing } - val fieldReference = resolveFieldNameFromExpression(filter.argumentList.expressions[0]) + val fieldReference = filter.argumentList.expressions[0].resolveFieldNameFromExpression() // if it's only 2 arguments it can be either: // - in(field, singleElement) -> valid because of varargs, becomes a single element array // - in(field, array) -> valid because of varargs @@ -343,7 +343,7 @@ object JavaDriverDialectParser : DialectParser { return null } val fieldExpression = filter.argumentList.expressions[0] - val fieldReference = resolveFieldNameFromExpression(fieldExpression) + val fieldReference = fieldExpression.resolveFieldNameFromExpression() val valueReference = HasValueReference.Inferred( source = fieldExpression, value = true, @@ -360,7 +360,7 @@ object JavaDriverDialectParser : DialectParser { ) } else if (method.parameters.size == 2) { // If it has two parameters, it's field/value. - val fieldReference = resolveFieldNameFromExpression(filter.argumentList.expressions[0]) + val fieldReference = filter.argumentList.expressions[0].resolveFieldNameFromExpression() val valueReference = resolveValueFromExpression(filter.argumentList.expressions[1]) return Node( @@ -393,7 +393,7 @@ object JavaDriverDialectParser : DialectParser { ) } else if (method.parameters.size == 2) { // If it has two parameters, it's field/value. - val fieldReference = resolveFieldNameFromExpression(filter.argumentList.expressions[0]) + val fieldReference = filter.argumentList.expressions[0].resolveFieldNameFromExpression() val valueReference = resolveValueFromExpression(filter.argumentList.expressions[1]) return Node( @@ -410,7 +410,7 @@ object JavaDriverDialectParser : DialectParser { ) } else if (method.parameters.size == 1) { // Updates.unset for example - val fieldReference = resolveFieldNameFromExpression(filter.argumentList.expressions[0]) + val fieldReference = filter.argumentList.expressions[0].resolveFieldNameFromExpression() return Node( filter, @@ -613,7 +613,7 @@ object JavaDriverDialectParser : DialectParser { "ascending", "descending" -> expression.getVarArgsOrIterableArgs() .mapNotNull { - val fieldReference = resolveFieldNameFromExpression(it) + val fieldReference = it.resolveFieldNameFromExpression() val methodName = Name.from(methodCall.name) when (fieldReference) { is FromSchema -> Node( @@ -818,16 +818,6 @@ object JavaDriverDialectParser : DialectParser { callMethod?.containingClass?.qualifiedName == AGGREGATES_FQN } - private fun resolveFieldNameFromExpression(expression: PsiExpression): HasFieldReference.FieldReference { - val fieldNameAsString = expression.tryToResolveAsConstantString() - val fieldReference = - fieldNameAsString?.let { - FromSchema(expression, it) - } ?: HasFieldReference.Unknown - - return fieldReference - } - private fun resolveValueFromExpression(expression: PsiExpression): HasValueReference.ValueReference { val (wasResolvedAtCompileTime, resolvedValue) = expression.tryToResolveAsConstant() @@ -907,6 +897,16 @@ object JavaDriverDialectParser : DialectParser { } } +fun PsiExpression.resolveFieldNameFromExpression(): HasFieldReference.FieldReference { + val fieldNameAsString = tryToResolveAsConstantString() + val fieldReference = + fieldNameAsString?.let { + FromSchema(this, it) + } ?: HasFieldReference.Unknown + + return fieldReference +} + fun PsiExpressionList.inferFromSingleArrayArgument(start: Int = 0): HasValueReference.ValueReference { val arrayArg = expressions[start] val (constant, value) = arrayArg.tryToResolveAsConstant() diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt new file mode 100644 index 00000000..d35a28ab --- /dev/null +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt @@ -0,0 +1,69 @@ +package com.mongodb.jbplugin.dialects.springcriteria + +import com.intellij.psi.PsiElement +import com.intellij.psi.PsiMethod +import com.intellij.psi.PsiMethodCallExpression +import com.mongodb.jbplugin.dialects.javadriver.glossary.fuzzyResolveMethod +import com.mongodb.jbplugin.dialects.javadriver.glossary.resolveToMethodCallExpression +import com.mongodb.jbplugin.dialects.springcriteria.aggregationparsers.MatchStageParser +import com.mongodb.jbplugin.mql.Node + +/** + * Parser for parsing supported patterns of writing an aggregation pipeline. + * Supported patterns for writing aggregation calls are the following: + * 1. MongoTemplate.aggregate() + * 2. MongoTemplate.aggregateStream() + * + * The AggregationParser concerns itself only with parsing the aggregation related semantics and + * leave the rest as a responsibility for the composing unit. + */ +class AggregationStagesParser(private val matchStageParser: MatchStageParser) { + private fun isStageCall(stageCallMethod: PsiMethod): Boolean { + return MatchStageParser.isMatchStageCall(stageCallMethod) + } + + private fun parseAggregationStages( + newAggregationCall: PsiMethodCallExpression + ): List> { + val newAggregationCallArguments = newAggregationCall.argumentList.expressions + val resolvedStageCalls = newAggregationCallArguments.mapNotNull { stageCallExpression -> + stageCallExpression.resolveToMethodCallExpression { _, stageCallMethod -> + isStageCall(stageCallMethod) + } + } + return resolvedStageCalls.map { stageCall -> + val stageCallMethod = stageCall.fuzzyResolveMethod() ?: return@map Node( + source = stageCall, + components = emptyList() + ) + + if (MatchStageParser.isMatchStageCall(stageCallMethod)) { + matchStageParser.parse(stageCall) + } else { + Node( + source = stageCall, + components = emptyList() + ) + } + } + } + + fun parse(aggregateRootCall: PsiMethodCallExpression): List> { + val aggregateRootCallArguments = aggregateRootCall.argumentList.expressions + + val newAggregationCallExpression = aggregateRootCallArguments.getOrNull(0) + ?: return emptyList() + + // This is the call to Aggregation.newAggregation method which is generally the first + // argument to the root aggregate call. All the aggregation stages are to be found as + // the argument to this method call. + val newAggregationCall = newAggregationCallExpression.resolveToMethodCallExpression { + _, + method + -> + method.name == "newAggregation" + } ?: return emptyList() + + return parseAggregationStages(newAggregationCall) + } +} diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt index d92b1d93..ae3841e3 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt @@ -10,6 +10,7 @@ import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtract import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtractor.extractCollectionFromQueryChain import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtractor.extractCollectionFromStringTypeParameter import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtractor.or +import com.mongodb.jbplugin.dialects.springcriteria.aggregationparsers.MatchStageParser import com.mongodb.jbplugin.mql.BsonAny import com.mongodb.jbplugin.mql.BsonArray import com.mongodb.jbplugin.mql.Node @@ -18,11 +19,8 @@ import com.mongodb.jbplugin.mql.toBsonType private const val CRITERIA_CLASS_FQN = "org.springframework.data.mongodb.core.query.Criteria" private const val DOCUMENT_FQN = "org.springframework.data.mongodb.core.mapping.Document" -private const val AGGREGATE_FQN = "org.springframework.data.mongodb.core.aggregation.Aggregation" private const val MONGO_TEMPLATE_FQN = "org.springframework.data.mongodb.core.MongoTemplate" -private val PARSEABLE_AGGREGATION_STAGE_METHODS = listOf( - "match" -) +const val AGGREGATE_FQN = "org.springframework.data.mongodb.core.aggregation.Aggregation" object SpringCriteriaDialectParser : DialectParser { override fun isCandidateForQuery(source: PsiElement) = @@ -218,19 +216,6 @@ object SpringCriteriaDialectParser : DialectParser { ) "aggregate", "aggregateStream" -> { val expressions = mongoOpCall.argumentList.expressions - val newAggregationCall = expressions.getOrNull(0)?.resolveToMethodCallExpression { - _, - method - -> - method.name == "newAggregation" - } - val resolvedStageCallExpression = - newAggregationCall?.getVarArgsOrIterableArgs()?.mapNotNull { - it.resolveToMethodCallExpression { _, method -> - method.containingClass?.qualifiedName == AGGREGATE_FQN && - PARSEABLE_AGGREGATION_STAGE_METHODS.contains(method.name) - } - } ?: emptyList() val collectionExpression = expressions.getOrNull(1) return Node( mongoOpCall, @@ -247,7 +232,9 @@ object SpringCriteriaDialectParser : DialectParser { extractCollectionFromStringTypeParameter(collectionExpression) ), HasAggregation( - parseAggregationStagesFromCurrentCall(resolvedStageCallExpression) + children = AggregationStagesParser( + matchStageParser = MatchStageParser(::parseFilterRecursively) + ).parse(mongoOpCall) ) ) ) @@ -271,7 +258,7 @@ object SpringCriteriaDialectParser : DialectParser { } override fun isReferenceToDatabase(source: PsiElement): Boolean { - return false // databases are in property files and we don't support AC there yet + return false // databases are in property files, and we don't support AC there yet } override fun isReferenceToCollection(source: PsiElement): Boolean { @@ -336,34 +323,6 @@ object SpringCriteriaDialectParser : DialectParser { ) } - private fun parseAggregationStagesFromCurrentCall( - stageCallExpressions: List - ): List> { - return stageCallExpressions.mapNotNull { stageCall -> - val stageMethod = stageCall.fuzzyResolveMethod() ?: return@mapNotNull Node( - source = stageCall, - components = emptyList() - ) - - when (stageMethod.name) { - "match" -> parseMatchStageCall(stageCall) - else -> null - } - } - } - - private fun parseMatchStageCall(matchStageCall: PsiMethodCallExpression): Node { - return Node( - source = matchStageCall, - components = listOf( - Named(Name.MATCH), - HasFilter( - parseFilterRecursively(matchStageCall.argumentList.expressions.getOrNull(0)) - ) - ) - ) - } - private fun parseFilterRecursively( valueFilterExpression: PsiElement? ): List> { diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt new file mode 100644 index 00000000..8fa668df --- /dev/null +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt @@ -0,0 +1,38 @@ +package com.mongodb.jbplugin.dialects.springcriteria.aggregationparsers + +import com.intellij.psi.PsiElement +import com.intellij.psi.PsiMethod +import com.intellij.psi.PsiMethodCallExpression +import com.mongodb.jbplugin.dialects.springcriteria.AGGREGATE_FQN +import com.mongodb.jbplugin.mql.Node +import com.mongodb.jbplugin.mql.components.HasFilter +import com.mongodb.jbplugin.mql.components.Name +import com.mongodb.jbplugin.mql.components.Named + +class MatchStageParser(private val parseFilters: (PsiElement) -> List>) { + private fun createMatchStageNode( + source: PsiElement, + filters: List> = emptyList() + ) = Node( + source = source, + components = listOf( + Named(Name.MATCH), + HasFilter(filters), + ) + ) + + fun parse(matchStageCall: PsiMethodCallExpression): Node { + val filterExpression = matchStageCall.argumentList.expressions.getOrNull(0) + ?: return createMatchStageNode(source = matchStageCall) + return createMatchStageNode( + source = matchStageCall, + filters = parseFilters(filterExpression) + ) + } + + companion object { + fun isMatchStageCall(method: PsiMethod): Boolean { + return method.containingClass?.qualifiedName != AGGREGATE_FQN && method.name == "match" + } + } +} diff --git a/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/AggregationParserTest.kt b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/AggregationStagesParserTest.kt similarity index 99% rename from packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/AggregationParserTest.kt rename to packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/AggregationStagesParserTest.kt index 7ca58924..59c567f7 100644 --- a/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/AggregationParserTest.kt +++ b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/AggregationStagesParserTest.kt @@ -16,7 +16,7 @@ import com.mongodb.jbplugin.mql.components.IsCommand import org.junit.jupiter.api.Assertions.assertEquals @IntegrationTest -class AggregationParserTest { +class AggregationStagesParserTest { @ParsingTest( fileName = "Book.java", """ From 4dee85aa7487bf316a257a599a8df6c80be39058 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 16 Jan 2025 12:34:33 +0100 Subject: [PATCH 2/9] chore: fix the bad condition for match stage call detection --- .../springcriteria/aggregationparsers/MatchStageParser.kt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt index 8fa668df..28abe750 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt @@ -32,7 +32,7 @@ class MatchStageParser(private val parseFilters: (PsiElement) -> List Date: Thu, 16 Jan 2025 13:17:15 +0100 Subject: [PATCH 3/9] chore: implement a StageParser interface --- .../springcriteria/AggregationStagesParser.kt | 6 ++--- .../SpringCriteriaDialectParser.kt | 2 +- .../MatchStageParser.kt | 25 ++++++++++--------- .../aggregationstageparsers/StageParser.kt | 11 ++++++++ .../AggregationStagesParserTest.kt | 9 +------ .../MatchStageParserTest.kt | 2 +- 6 files changed, 30 insertions(+), 25 deletions(-) rename packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/{aggregationparsers => aggregationstageparsers}/MatchStageParser.kt (59%) create mode 100644 packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/StageParser.kt rename packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/{aggregationparser => }/AggregationStagesParserTest.kt (94%) rename packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/{aggregationparser => aggregationstageparsers}/MatchStageParserTest.kt (99%) diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt index d35a28ab..db9ab95f 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt @@ -5,7 +5,7 @@ import com.intellij.psi.PsiMethod import com.intellij.psi.PsiMethodCallExpression import com.mongodb.jbplugin.dialects.javadriver.glossary.fuzzyResolveMethod import com.mongodb.jbplugin.dialects.javadriver.glossary.resolveToMethodCallExpression -import com.mongodb.jbplugin.dialects.springcriteria.aggregationparsers.MatchStageParser +import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.MatchStageParser import com.mongodb.jbplugin.mql.Node /** @@ -19,7 +19,7 @@ import com.mongodb.jbplugin.mql.Node */ class AggregationStagesParser(private val matchStageParser: MatchStageParser) { private fun isStageCall(stageCallMethod: PsiMethod): Boolean { - return MatchStageParser.isMatchStageCall(stageCallMethod) + return matchStageParser.canParse(stageCallMethod) } private fun parseAggregationStages( @@ -37,7 +37,7 @@ class AggregationStagesParser(private val matchStageParser: MatchStageParser) { components = emptyList() ) - if (MatchStageParser.isMatchStageCall(stageCallMethod)) { + if (matchStageParser.canParse(stageCallMethod)) { matchStageParser.parse(stageCall) } else { Node( diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt index ae3841e3..1241315b 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt @@ -10,7 +10,7 @@ import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtract import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtractor.extractCollectionFromQueryChain import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtractor.extractCollectionFromStringTypeParameter import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtractor.or -import com.mongodb.jbplugin.dialects.springcriteria.aggregationparsers.MatchStageParser +import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.MatchStageParser import com.mongodb.jbplugin.mql.BsonAny import com.mongodb.jbplugin.mql.BsonArray import com.mongodb.jbplugin.mql.Node diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParser.kt similarity index 59% rename from packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt rename to packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParser.kt index 28abe750..f81b1e36 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparsers/MatchStageParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParser.kt @@ -1,4 +1,4 @@ -package com.mongodb.jbplugin.dialects.springcriteria.aggregationparsers +package com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers import com.intellij.psi.PsiElement import com.intellij.psi.PsiMethod @@ -9,7 +9,9 @@ import com.mongodb.jbplugin.mql.components.HasFilter import com.mongodb.jbplugin.mql.components.Name import com.mongodb.jbplugin.mql.components.Named -class MatchStageParser(private val parseFilters: (PsiElement) -> List>) { +class MatchStageParser( + private val parseFilters: (PsiElement) -> List> +) : StageParser { private fun createMatchStageNode( source: PsiElement, filters: List> = emptyList() @@ -21,18 +23,17 @@ class MatchStageParser(private val parseFilters: (PsiElement) -> List { - val filterExpression = matchStageCall.argumentList.expressions.getOrNull(0) - ?: return createMatchStageNode(source = matchStageCall) + override fun canParse(stageCallMethod: PsiMethod): Boolean { + return stageCallMethod.containingClass?.qualifiedName == AGGREGATE_FQN && + stageCallMethod.name == "match" + } + + override fun parse(stageCall: PsiMethodCallExpression): Node { + val filterExpression = stageCall.argumentList.expressions.getOrNull(0) + ?: return createMatchStageNode(source = stageCall) return createMatchStageNode( - source = matchStageCall, + source = stageCall, filters = parseFilters(filterExpression) ) } - - companion object { - fun isMatchStageCall(method: PsiMethod): Boolean { - return method.containingClass?.qualifiedName == AGGREGATE_FQN && method.name == "match" - } - } } diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/StageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/StageParser.kt new file mode 100644 index 00000000..e0f504cc --- /dev/null +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/StageParser.kt @@ -0,0 +1,11 @@ +package com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers + +import com.intellij.psi.PsiElement +import com.intellij.psi.PsiMethod +import com.intellij.psi.PsiMethodCallExpression +import com.mongodb.jbplugin.mql.Node + +interface StageParser { + fun canParse(stageCallMethod: PsiMethod): Boolean + fun parse(stageCall: PsiMethodCallExpression): Node +} diff --git a/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/AggregationStagesParserTest.kt b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParserTest.kt similarity index 94% rename from packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/AggregationStagesParserTest.kt rename to packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParserTest.kt index 59c567f7..c345f32b 100644 --- a/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/AggregationStagesParserTest.kt +++ b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParserTest.kt @@ -1,14 +1,7 @@ -package com.mongodb.jbplugin.dialects.springcriteria.aggregationparser +package com.mongodb.jbplugin.dialects.springcriteria import com.intellij.psi.PsiElement import com.intellij.psi.PsiFile -import com.mongodb.jbplugin.dialects.springcriteria.IntegrationTest -import com.mongodb.jbplugin.dialects.springcriteria.ParsingTest -import com.mongodb.jbplugin.dialects.springcriteria.SpringCriteriaDialectParser -import com.mongodb.jbplugin.dialects.springcriteria.assert -import com.mongodb.jbplugin.dialects.springcriteria.collection -import com.mongodb.jbplugin.dialects.springcriteria.component -import com.mongodb.jbplugin.dialects.springcriteria.getQueryAtMethod import com.mongodb.jbplugin.mql.components.HasAggregation import com.mongodb.jbplugin.mql.components.HasCollectionReference import com.mongodb.jbplugin.mql.components.HasSourceDialect diff --git a/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/MatchStageParserTest.kt b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParserTest.kt similarity index 99% rename from packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/MatchStageParserTest.kt rename to packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParserTest.kt index d5348193..a173bf99 100644 --- a/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationparser/MatchStageParserTest.kt +++ b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParserTest.kt @@ -1,4 +1,4 @@ -package com.mongodb.jbplugin.dialects.springcriteria.aggregationparser +package com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers import com.intellij.psi.PsiElement import com.intellij.psi.PsiFile From 1e0ca1ad563024017fd14efea7fe87abbdd242cf Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 16 Jan 2025 14:29:49 +0100 Subject: [PATCH 4/9] feat: adds support for parsing Aggregate.project includes support for parsing chained calls(andInclude, andExclude) --- .../SpringCriteriaRepository.java | 2 +- .../springcriteria/AggregationStagesParser.kt | 8 +- .../SpringCriteriaDialectParser.kt | 2 + .../ProjectStageParser.kt | 229 +++++ .../springcriteria/IntegrationTest.kt | 25 + .../ProjectStageParserTest.kt | 838 ++++++++++++++++++ 6 files changed, 1102 insertions(+), 2 deletions(-) create mode 100644 packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt create mode 100644 packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParserTest.kt diff --git a/packages/jetbrains-plugin/src/test/resources/project-fixtures/basic-java-project-with-mongodb/src/main/java/alt/mongodb/springcriteria/SpringCriteriaRepository.java b/packages/jetbrains-plugin/src/test/resources/project-fixtures/basic-java-project-with-mongodb/src/main/java/alt/mongodb/springcriteria/SpringCriteriaRepository.java index d3229d3b..1f4fbd44 100644 --- a/packages/jetbrains-plugin/src/test/resources/project-fixtures/basic-java-project-with-mongodb/src/main/java/alt/mongodb/springcriteria/SpringCriteriaRepository.java +++ b/packages/jetbrains-plugin/src/test/resources/project-fixtures/basic-java-project-with-mongodb/src/main/java/alt/mongodb/springcriteria/SpringCriteriaRepository.java @@ -35,7 +35,7 @@ private List allMoviesWithRatingAtLeastAgg(int rating) { return template.aggregate( Aggregation.newAggregation( Aggregation.match(where( "tomatoes.viewer.rating").gte(rating)), - Aggregation.project("asd").andInclude("asd").andExclude("qwe") + Aggregation.project("fieldA").andInclude("fieldB").andExclude("fieldC") ), Movie.class, Movie.class diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt index db9ab95f..266d9b77 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt @@ -6,6 +6,7 @@ import com.intellij.psi.PsiMethodCallExpression import com.mongodb.jbplugin.dialects.javadriver.glossary.fuzzyResolveMethod import com.mongodb.jbplugin.dialects.javadriver.glossary.resolveToMethodCallExpression import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.MatchStageParser +import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.ProjectStageParser import com.mongodb.jbplugin.mql.Node /** @@ -18,8 +19,11 @@ import com.mongodb.jbplugin.mql.Node * leave the rest as a responsibility for the composing unit. */ class AggregationStagesParser(private val matchStageParser: MatchStageParser) { + private val projectStageParser = ProjectStageParser() + private fun isStageCall(stageCallMethod: PsiMethod): Boolean { - return matchStageParser.canParse(stageCallMethod) + return matchStageParser.canParse(stageCallMethod) || + projectStageParser.canParse(stageCallMethod) } private fun parseAggregationStages( @@ -39,6 +43,8 @@ class AggregationStagesParser(private val matchStageParser: MatchStageParser) { if (matchStageParser.canParse(stageCallMethod)) { matchStageParser.parse(stageCall) + } else if (projectStageParser.canParse(stageCallMethod)) { + projectStageParser.parse(stageCall) } else { Node( source = stageCall, diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt index 1241315b..c791dd8e 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt @@ -21,6 +21,8 @@ private const val CRITERIA_CLASS_FQN = "org.springframework.data.mongodb.core.qu private const val DOCUMENT_FQN = "org.springframework.data.mongodb.core.mapping.Document" private const val MONGO_TEMPLATE_FQN = "org.springframework.data.mongodb.core.MongoTemplate" const val AGGREGATE_FQN = "org.springframework.data.mongodb.core.aggregation.Aggregation" +const val PROJECTION_OPERATION_FQN = "org.springframework.data.mongodb.core.aggregation.ProjectionOperation" +const val FIELDS_FQN = "org.springframework.data.mongodb.core.aggregation.Fields" object SpringCriteriaDialectParser : DialectParser { override fun isCandidateForQuery(source: PsiElement) = diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt new file mode 100644 index 00000000..7d913641 --- /dev/null +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt @@ -0,0 +1,229 @@ +package com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers + +import com.intellij.psi.PsiElement +import com.intellij.psi.PsiExpression +import com.intellij.psi.PsiMethod +import com.intellij.psi.PsiMethodCallExpression +import com.mongodb.jbplugin.dialects.javadriver.glossary.fuzzyResolveMethod +import com.mongodb.jbplugin.dialects.javadriver.glossary.resolveFieldNameFromExpression +import com.mongodb.jbplugin.dialects.javadriver.glossary.resolveToMethodCallExpression +import com.mongodb.jbplugin.dialects.springcriteria.AGGREGATE_FQN +import com.mongodb.jbplugin.dialects.springcriteria.FIELDS_FQN +import com.mongodb.jbplugin.dialects.springcriteria.PROJECTION_OPERATION_FQN +import com.mongodb.jbplugin.mql.BsonInt32 +import com.mongodb.jbplugin.mql.Node +import com.mongodb.jbplugin.mql.components.HasFieldReference +import com.mongodb.jbplugin.mql.components.HasProjections +import com.mongodb.jbplugin.mql.components.HasValueReference +import com.mongodb.jbplugin.mql.components.Name +import com.mongodb.jbplugin.mql.components.Named + +class ProjectStageParser : StageParser { + private fun createProjectNode( + source: PsiElement, + projections: List> + ): Node { + return Node( + source = source, + components = listOf( + Named(Name.PROJECT), + HasProjections( + children = projections + ) + ) + ) + } + + private fun createProjectedFieldNode( + fieldExpression: PsiExpression, + projectionName: Name + ): Node { + val fieldReference = fieldExpression.resolveFieldNameFromExpression() + return Node( + source = fieldExpression as PsiElement, + components = listOf( + Named(projectionName), + HasFieldReference(fieldReference), + HasValueReference( + reference = when (projectionName) { + Name.INCLUDE, + Name.EXCLUDE -> HasValueReference.Inferred( + source = fieldExpression, + value = if (projectionName == Name.INCLUDE) 1 else 0, + type = BsonInt32 + ) + else -> HasValueReference.Unknown + } + ) + ) + ) + } + + /** + * Parses a `PsiExpression` that represents a `Field` object and returns a `Node` for the field + * referenced in the `Field` object. + * + * The only helper that we support for creating a `Field` object is `Fields.field()` with the + * following alternative signatures: + * 1. `Fields.field("fieldA")` + * 2. `Fields.field(fieldBFromVariable)` + * 3. `Fields.field(fieldCFromMethodCall())` + */ + private fun fieldObjectExpressionToProjectedFieldNode( + fieldObjectExpression: PsiExpression, + projectionName: Name + ): Node? { + val resolvedFieldMethodCall = fieldObjectExpression.resolveToMethodCallExpression { + _, + fieldObjectMethod + -> + fieldObjectMethod.containingClass?.qualifiedName == FIELDS_FQN && + fieldObjectMethod.name == "field" + } ?: return null + + val resolvedFieldMethod = resolvedFieldMethodCall.fuzzyResolveMethod() ?: return null + // This represents the following call: + // `Fields.field("fieldAlias", "actualFieldInDocument")` + // and since this translates to a rename operation, which we do not support just yet, we + // will ignore this until we come back to this. + if (resolvedFieldMethodCall.argumentList.expressions.size == 2) { + return null + } else { + val fieldExpression = resolvedFieldMethodCall.argumentList.expressions.getOrNull(0) + ?: return null + return createProjectedFieldNode(fieldExpression, projectionName) + } + } + + /** + * Parses a `PsiExpression` that represents a `Fields` object and returns a `Node` for each + * field referenced in the `Fields` object. + * + * The two different ways of generating a `Fields` object that we support and parse in here are: + * 1. `Fields.fields("fieldA", fieldBFromVariable, fieldCFromMethodCall())` + * 2. `Fields.from(Fields.field("fieldA"), Fields.field(fieldBFromVariable))` + */ + private fun fieldsObjectExpressionToProjectedFieldNode( + fieldsObjectExpression: PsiExpression, + projectionName: Name + ): List> { + val resolvedFieldsMethodCall = fieldsObjectExpression.resolveToMethodCallExpression { + _, + fieldsObjectMethod + -> + fieldsObjectMethod.containingClass?.qualifiedName == FIELDS_FQN && + (fieldsObjectMethod.name == "fields" || fieldsObjectMethod.name == "from") + } ?: return emptyList() + + val resolvedFieldsMethod = resolvedFieldsMethodCall.fuzzyResolveMethod() + ?: return emptyList() + + return when (resolvedFieldsMethod.name) { + "fields" -> resolvedFieldsMethodCall.argumentList.expressions.mapNotNull { + createProjectedFieldNode(it, projectionName) + } + "from" -> resolvedFieldsMethodCall.argumentList.expressions.mapNotNull { + fieldObjectExpressionToProjectedFieldNode(it, projectionName) + } + else -> { + // log.warn("Unsupported Fields method: ${resolvedFieldsMethod.name}") + emptyList() + } + } + } + + /** + * Parses a method that has the following overloaded signatures: + * 1. methodCall(String... fieldNames) + * 2. methodCall(Fields fields) + */ + private fun parseMethodCallWithStringVarArgsAndFields( + methodCall: PsiMethodCallExpression, + projectionName: Name, + ): List> { + val method = methodCall.fuzzyResolveMethod() ?: return emptyList() + return if (method.isVarArgs) { + methodCall.argumentList.expressions.mapNotNull { + createProjectedFieldNode(fieldExpression = it, projectionName = projectionName) + } + } else { + val maybeFieldsObjectExpression = methodCall.argumentList.expressions.getOrNull(0) + ?: return emptyList() + fieldsObjectExpressionToProjectedFieldNode(maybeFieldsObjectExpression, projectionName) + } + } + + private fun parseAndIncludeCall(methodCall: PsiMethodCallExpression): List> { + // andInclude have two overloads, one which takes string varargs and another + // which takes a Fields object + return parseMethodCallWithStringVarArgsAndFields(methodCall, Name.INCLUDE) + } + + private fun parseAndExcludeCall(methodCall: PsiMethodCallExpression): List> { + // andExclude only accepts a varargs of strings as arguments + return methodCall.argumentList.expressions.mapNotNull { + createProjectedFieldNode(fieldExpression = it, projectionName = Name.EXCLUDE) + } + } + + private fun parseProjectCall(methodCall: PsiMethodCallExpression): List> { + // Aggregation.project() has three overloads + // 1. Aggregation.project(String... fieldNames) + // 2. Aggregation.project(Fields fields) + // 3. Aggregation.project(Class type) + // We only support the first 2 and ignore the third one + return parseMethodCallWithStringVarArgsAndFields(methodCall, Name.INCLUDE) + } + + override fun canParse(stageCallMethod: PsiMethod): Boolean { + val methodFqn = stageCallMethod.containingClass?.qualifiedName ?: return false + // Project stage call might contain chained operations which might result + // in a method call from `PROJECTION_OPERATION_FQN` so we account for both + return listOf(AGGREGATE_FQN, PROJECTION_OPERATION_FQN).contains(methodFqn) && + // we consider the root call project and any chained calls which we currently + // support + listOf("project", "andInclude", "andExclude").contains(stageCallMethod.name) + } + + override fun parse(stageCall: PsiMethodCallExpression): Node { + return createProjectNode( + source = stageCall, + projections = stageCall.gatherChainedCalls().flatMap { methodCall -> + val method = methodCall.fuzzyResolveMethod() ?: return@flatMap emptyList() + when (method.name) { + "andInclude" -> parseAndIncludeCall(methodCall) + "andExclude" -> parseAndExcludeCall(methodCall) + "project" -> parseProjectCall(methodCall) + else -> emptyList() + } + } + ) + } +} + +/** + * From a PsiMethodCallExpression, it attempts to travel upwards the chain(assuming there is one) + * while gathering other PsiMethodCallExpressions it comes across, until there is no further + * PsiMethodCallExpression in the chain. For example, consider the following method call + * ``` + * Aggregation.project().andInclude("fieldA").andExclude("_id") + * ``` + * and given that the method was called for PsiMethodCallExpression referring the entire call above, + * the chain will consist of the following method calls + * [ + * Aggregation.project().andInclude("fieldA").andExclude("_id"), + * Aggregation.project().andInclude("fieldA") + * Aggregation.project() + * ] + */ +fun PsiMethodCallExpression.gatherChainedCalls(): List { + val chain = mutableListOf() + var currentCall: PsiMethodCallExpression? = this + + while (currentCall != null) { + chain.add(currentCall) + currentCall = currentCall.methodExpression.qualifierExpression as? PsiMethodCallExpression + } + + return chain +} diff --git a/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/IntegrationTest.kt b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/IntegrationTest.kt index c91ea332..888f252a 100644 --- a/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/IntegrationTest.kt +++ b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/IntegrationTest.kt @@ -39,6 +39,7 @@ import com.mongodb.jbplugin.mql.components.HasAggregation import com.mongodb.jbplugin.mql.components.HasCollectionReference import com.mongodb.jbplugin.mql.components.HasFieldReference import com.mongodb.jbplugin.mql.components.HasFilter +import com.mongodb.jbplugin.mql.components.HasProjections import com.mongodb.jbplugin.mql.components.HasUpdates import com.mongodb.jbplugin.mql.components.HasValueReference import com.mongodb.jbplugin.mql.components.IsCommand @@ -352,6 +353,30 @@ fun Node.filterN( filter.assertions() } +fun Node.projectionN( + n: Int, + name: Name? = null, + assertions: Node.() -> Unit = { + } +) { + val projections = component>() + assertNotNull(projections) + + val projection = projections!!.children[n] + + if (name != null) { + val qname = projection.component() + assertNotEquals(null, qname) { + "Expected a named operation with name $name but null found." + } + assertEquals(name, qname?.name) { + "Expected a named operation with name $name but $qname found." + } + } + + projection.assertions() +} + fun Node.updateN( n: Int, name: Name? = null, diff --git a/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParserTest.kt b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParserTest.kt new file mode 100644 index 00000000..dce50de7 --- /dev/null +++ b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParserTest.kt @@ -0,0 +1,838 @@ +package com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers + +import com.intellij.psi.PsiElement +import com.intellij.psi.PsiFile +import com.mongodb.jbplugin.dialects.springcriteria.IntegrationTest +import com.mongodb.jbplugin.dialects.springcriteria.ParsingTest +import com.mongodb.jbplugin.dialects.springcriteria.SpringCriteriaDialectParser +import com.mongodb.jbplugin.dialects.springcriteria.assert +import com.mongodb.jbplugin.dialects.springcriteria.collection +import com.mongodb.jbplugin.dialects.springcriteria.component +import com.mongodb.jbplugin.dialects.springcriteria.field +import com.mongodb.jbplugin.dialects.springcriteria.getQueryAtMethod +import com.mongodb.jbplugin.dialects.springcriteria.projectionN +import com.mongodb.jbplugin.dialects.springcriteria.stageN +import com.mongodb.jbplugin.dialects.springcriteria.value +import com.mongodb.jbplugin.mql.BsonInt32 +import com.mongodb.jbplugin.mql.components.HasCollectionReference +import com.mongodb.jbplugin.mql.components.HasFieldReference +import com.mongodb.jbplugin.mql.components.HasProjections +import com.mongodb.jbplugin.mql.components.HasSourceDialect +import com.mongodb.jbplugin.mql.components.HasValueReference +import com.mongodb.jbplugin.mql.components.IsCommand +import com.mongodb.jbplugin.mql.components.Name +import org.junit.jupiter.api.Assertions.assertEquals + +@IntegrationTest +class ProjectStageParserTest { + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public AggregationResults allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project() + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse an empty project stage`(psiFile: PsiFile) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(0, children.size) + } + } + } + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public AggregationResults allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude().andExclude() + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse an empty project stage with empty chained calls`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(0, children.size) + } + } + } + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + private String fieldCFromMethodCall() { + return "fieldC"; + } + + public AggregationResults allReleasedBooks() { + String fieldBFromVariable = "fieldB"; + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project("fieldA", fieldBFromVariable, fieldCFromMethodCall()) + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse string field names passed as varargs to Aggregation#project`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(3, children.size) + } + + projectionN(0, Name.INCLUDE) { + field> { + assertEquals("fieldA", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(1, Name.INCLUDE) { + field> { + assertEquals("fieldB", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(2, Name.INCLUDE) { + field> { + assertEquals("fieldC", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + } + } + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + private String fieldCFromMethodCall() { + return "fieldC"; + } + + public AggregationResults allReleasedBooks() { + String fieldBFromVariable = "fieldB"; + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project( + Fields.fields("fieldA", fieldBFromVariable, fieldCFromMethodCall()) + ) + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse field names provided via Fields#fields to Aggregation#project`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(3, children.size) + } + + projectionN(0, Name.INCLUDE) { + field> { + assertEquals("fieldA", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(1, Name.INCLUDE) { + field> { + assertEquals("fieldB", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(2, Name.INCLUDE) { + field> { + assertEquals("fieldC", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + } + } + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + private String fieldCFromMethodCall() { + return "fieldC"; + } + + public AggregationResults allReleasedBooks() { + String fieldBFromVariable = "fieldB"; + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project( + Fields.from( + Fields.field("fieldA"), + Fields.field(fieldBFromVariable), + Fields.field(fieldCFromMethodCall()) + ) + ) + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse field names provided via Field#field objects to Aggregation#project`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(3, children.size) + } + + projectionN(0, Name.INCLUDE) { + field> { + assertEquals("fieldA", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(1, Name.INCLUDE) { + field> { + assertEquals("fieldB", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(2, Name.INCLUDE) { + field> { + assertEquals("fieldC", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + } + } + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + private String fieldCFromMethodCall() { + return "fieldC"; + } + + public AggregationResults allReleasedBooks() { + String fieldBFromVariable = "fieldB"; + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude("fieldA", fieldBFromVariable, fieldCFromMethodCall()) + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse chained andInclude calls with string field names passed as varargs`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(3, children.size) + } + + projectionN(0, Name.INCLUDE) { + field> { + assertEquals("fieldA", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(1, Name.INCLUDE) { + field> { + assertEquals("fieldB", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(2, Name.INCLUDE) { + field> { + assertEquals("fieldC", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + } + } + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + private String fieldCFromMethodCall() { + return "fieldC"; + } + + public AggregationResults allReleasedBooks() { + String fieldBFromVariable = "fieldB"; + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude(Fields.fields("fieldA", fieldBFromVariable, fieldCFromMethodCall())) + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse chained andInclude calls with field names provided via Fields#fields`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(3, children.size) + } + + projectionN(0, Name.INCLUDE) { + field> { + assertEquals("fieldA", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(1, Name.INCLUDE) { + field> { + assertEquals("fieldB", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(2, Name.INCLUDE) { + field> { + assertEquals("fieldC", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + } + } + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + private String fieldCFromMethodCall() { + return "fieldC"; + } + + public AggregationResults allReleasedBooks() { + String fieldBFromVariable = "fieldB"; + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude( + Fields.from( + Fields.field("fieldA"), + Fields.field(fieldBFromVariable), + Fields.field(fieldCFromMethodCall()) + ) + ) + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse chained andInclude calls with field names provided via Fields#Field objects`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(3, children.size) + } + + projectionN(0, Name.INCLUDE) { + field> { + assertEquals("fieldA", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(1, Name.INCLUDE) { + field> { + assertEquals("fieldB", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(2, Name.INCLUDE) { + field> { + assertEquals("fieldC", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + } + } + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + private String fieldCFromMethodCall() { + return "fieldC"; + } + + public AggregationResults allReleasedBooks() { + String fieldBFromVariable = "fieldB"; + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andExclude("fieldA", fieldBFromVariable, fieldCFromMethodCall()) + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse chained andExclude calls with string field names passed as varargs`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(3, children.size) + } + + projectionN(0, Name.EXCLUDE) { + field> { + assertEquals("fieldA", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(0, value) + } + } + + projectionN(1, Name.EXCLUDE) { + field> { + assertEquals("fieldB", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(0, value) + } + } + + projectionN(2, Name.EXCLUDE) { + field> { + assertEquals("fieldC", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(0, value) + } + } + } + } + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; + +import java.util.List; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + private String fieldCFromMethodCall() { + return "fieldC"; + } + + public AggregationResults allReleasedBooks() { + String fieldBFromVariable = "fieldB"; + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project("fieldA").andInclude(fieldBFromVariable).andExclude(fieldCFromMethodCall()) + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse multiple chained calls`(psiFile: PsiFile) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + stageN(0, Name.PROJECT) { + component> { + assertEquals(3, children.size) + } + + projectionN(0, Name.EXCLUDE) { + field> { + assertEquals("fieldC", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(0, value) + } + } + + projectionN(1, Name.INCLUDE) { + field> { + assertEquals("fieldB", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + + projectionN(2, Name.INCLUDE) { + field> { + assertEquals("fieldA", fieldName) + } + value> { + assertEquals(BsonInt32, type) + assertEquals(1, value) + } + } + } + } + } +} From 4614a33176d46dc0f5e0aaafd12f342eb773400a Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 16 Jan 2025 14:49:00 +0100 Subject: [PATCH 5/9] chore: refactor AggregationStagesParser interface to accept a list of parsers --- .../springcriteria/AggregationStagesParser.kt | 25 ++++++------------- .../SpringCriteriaDialectParser.kt | 6 ++++- 2 files changed, 13 insertions(+), 18 deletions(-) diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt index 266d9b77..53d72807 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/AggregationStagesParser.kt @@ -5,8 +5,7 @@ import com.intellij.psi.PsiMethod import com.intellij.psi.PsiMethodCallExpression import com.mongodb.jbplugin.dialects.javadriver.glossary.fuzzyResolveMethod import com.mongodb.jbplugin.dialects.javadriver.glossary.resolveToMethodCallExpression -import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.MatchStageParser -import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.ProjectStageParser +import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.StageParser import com.mongodb.jbplugin.mql.Node /** @@ -18,12 +17,9 @@ import com.mongodb.jbplugin.mql.Node * The AggregationParser concerns itself only with parsing the aggregation related semantics and * leave the rest as a responsibility for the composing unit. */ -class AggregationStagesParser(private val matchStageParser: MatchStageParser) { - private val projectStageParser = ProjectStageParser() - +class AggregationStagesParser(private val stageParsers: List) { private fun isStageCall(stageCallMethod: PsiMethod): Boolean { - return matchStageParser.canParse(stageCallMethod) || - projectStageParser.canParse(stageCallMethod) + return stageParsers.any { it.canParse(stageCallMethod) } } private fun parseAggregationStages( @@ -41,16 +37,11 @@ class AggregationStagesParser(private val matchStageParser: MatchStageParser) { components = emptyList() ) - if (matchStageParser.canParse(stageCallMethod)) { - matchStageParser.parse(stageCall) - } else if (projectStageParser.canParse(stageCallMethod)) { - projectStageParser.parse(stageCall) - } else { - Node( - source = stageCall, - components = emptyList() - ) - } + val parser = stageParsers.find { it.canParse(stageCallMethod) } + parser?.parse(stageCall) ?: Node( + source = stageCall, + components = emptyList() + ) } } diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt index c791dd8e..cdf96d11 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt @@ -11,6 +11,7 @@ import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtract import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtractor.extractCollectionFromStringTypeParameter import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtractor.or import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.MatchStageParser +import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.ProjectStageParser import com.mongodb.jbplugin.mql.BsonAny import com.mongodb.jbplugin.mql.BsonArray import com.mongodb.jbplugin.mql.Node @@ -235,7 +236,10 @@ object SpringCriteriaDialectParser : DialectParser { ), HasAggregation( children = AggregationStagesParser( - matchStageParser = MatchStageParser(::parseFilterRecursively) + stageParsers = listOf( + MatchStageParser(::parseFilterRecursively), + ProjectStageParser() + ) ).parse(mongoOpCall) ) ) From 3aa1f1df39994a3e47762f7907d2cf72f9836c15 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 16 Jan 2025 15:07:54 +0100 Subject: [PATCH 6/9] chore: added tests for field existence inspection --- ...ngCriteriaFieldCheckLinterInspectionTest.kt | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/inspections/impl/SpringCriteriaFieldCheckLinterInspectionTest.kt b/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/inspections/impl/SpringCriteriaFieldCheckLinterInspectionTest.kt index 82a1d81b..9490f8f8 100644 --- a/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/inspections/impl/SpringCriteriaFieldCheckLinterInspectionTest.kt +++ b/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/inspections/impl/SpringCriteriaFieldCheckLinterInspectionTest.kt @@ -190,11 +190,29 @@ class BookRepository { ); } + String releasedFromMethodCall() { + return "released"; + } + public void allReleasedBooksAggregate() { + String releasedAsVariable = "released"; template.aggregate( Aggregation.newAggregation( Aggregation.match( where("released").is(true) + ), + Aggregation.project( + "released", + releasedAsVariable, + releasedFromMethodCall() + ).andInclude( + "released", + releasedAsVariable, + releasedFromMethodCall() + ).andExclude( + "released", + releasedAsVariable, + releasedFromMethodCall() ) ), Book.class, From 32f4cb866bd525b7e9dd09380b597967250cb514 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 16 Jan 2025 16:18:11 +0100 Subject: [PATCH 7/9] feat: AC support in project and chained calls --- ...SpringCriteriaCompletionContributorTest.kt | 518 +++++++++++++++++ ...iaMongoDbAutocompletionPopupHandlerTest.kt | 525 ++++++++++++++++++ .../SpringCriteriaDialectParser.kt | 24 +- .../MatchStageParser.kt | 8 + .../ProjectStageParser.kt | 17 + .../aggregationstageparsers/StageParser.kt | 4 + 6 files changed, 1091 insertions(+), 5 deletions(-) diff --git a/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/autocomplete/SpringCriteriaCompletionContributorTest.kt b/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/autocomplete/SpringCriteriaCompletionContributorTest.kt index b7f3c38b..302f6e64 100644 --- a/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/autocomplete/SpringCriteriaCompletionContributorTest.kt +++ b/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/autocomplete/SpringCriteriaCompletionContributorTest.kt @@ -392,4 +392,522 @@ class Repository { }, ) } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project("") + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in a Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project(Fields.fields("")) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in fields helper passed to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project(Fields.from(Fields.field(""))) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in field helper passed to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude("") + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in andInclude call chained to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude(Fields.fields("")) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in fields helper passed to andInclude call chained to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude(Fields.from(Fields.field(""))) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in field helper passed to andInclude call chained to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andExclude("") + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in andExclude call chained to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } } diff --git a/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/autocomplete/SpringCriteriaMongoDbAutocompletionPopupHandlerTest.kt b/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/autocomplete/SpringCriteriaMongoDbAutocompletionPopupHandlerTest.kt index 7e5a9e1f..f03e366e 100644 --- a/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/autocomplete/SpringCriteriaMongoDbAutocompletionPopupHandlerTest.kt +++ b/packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/autocomplete/SpringCriteriaMongoDbAutocompletionPopupHandlerTest.kt @@ -399,4 +399,529 @@ class Repository { }, ) } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project() + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in a Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + fixture.type('"') + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project(Fields.fields()) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in fields helper passed to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + fixture.type('"') + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project(Fields.from(Fields.field())) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in field helper passed to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + fixture.type('"') + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude() + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in andInclude call chained to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + fixture.type('"') + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude(Fields.fields()) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in fields helper passed to andInclude call chained to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + fixture.type('"') + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andInclude(Fields.from(Fields.field())) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in field helper passed to andInclude call chained to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + fixture.type('"') + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; +import org.springframework.data.mongodb.core.query.Criteria; + +import java.util.List; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + public List allReleasedBooks() { + return template.aggregate( + Aggregation.newAggregation( + Aggregation.project().andExclude() + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in andExclude call chained to Aggregation#project call`( + fixture: CodeInsightTestFixture, + ) { + val (dataSource, readModelProvider) = fixture.setupConnection() + fixture.specifyDatabase("myDatabase") + fixture.specifyDialect(SpringCriteriaDialect) + + val namespace = Namespace("myDatabase", "book") + + `when`( + readModelProvider.slice(eq(dataSource), eq(GetCollectionSchema.Slice(namespace))) + ).thenReturn( + GetCollectionSchema( + CollectionSchema( + namespace, + BsonObject( + mapOf( + "myField" to BsonString, + "myField2" to BsonString, + ), + ), + ), + ), + ) + + fixture.type('"') + val elements = fixture.completeBasic() + + assertTrue( + elements.containsElements { + it.lookupString == "myField" + }, + ) + + assertTrue( + elements.containsElements { + it.lookupString == "myField2" + }, + ) + } } diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt index cdf96d11..e5a41b99 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/SpringCriteriaDialectParser.kt @@ -12,6 +12,7 @@ import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtract import com.mongodb.jbplugin.dialects.springcriteria.QueryTargetCollectionExtractor.or import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.MatchStageParser import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.ProjectStageParser +import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.StageParser import com.mongodb.jbplugin.mql.BsonAny import com.mongodb.jbplugin.mql.BsonArray import com.mongodb.jbplugin.mql.Node @@ -26,6 +27,11 @@ const val PROJECTION_OPERATION_FQN = "org.springframework.data.mongodb.core.aggr const val FIELDS_FQN = "org.springframework.data.mongodb.core.aggregation.Fields" object SpringCriteriaDialectParser : DialectParser { + private val aggregationStageParsers: List = listOf( + MatchStageParser(::parseFilterRecursively), + ProjectStageParser() + ) + override fun isCandidateForQuery(source: PsiElement) = inferCommandFromMethod((source as? PsiMethodCallExpression)?.fuzzyResolveMethod()).type != IsCommand.CommandType.UNKNOWN @@ -236,10 +242,7 @@ object SpringCriteriaDialectParser : DialectParser { ), HasAggregation( children = AggregationStagesParser( - stageParsers = listOf( - MatchStageParser(::parseFilterRecursively), - ProjectStageParser() - ) + stageParsers = aggregationStageParsers ).parse(mongoOpCall) ) ) @@ -295,7 +298,11 @@ object SpringCriteriaDialectParser : DialectParser { } } - return isString && methodCall.isCriteriaExpression() + return isString && + ( + methodCall.isCriteriaExpression() || + methodCall.isSuitableForFieldAutoCompleteInAggregation(aggregationStageParsers) + ) } private fun isInsideDocAnnotations(source: PsiElement): Boolean { @@ -602,6 +609,13 @@ fun PsiMethodCallExpression.isCriteriaExpression(): Boolean { return method.containingClass?.qualifiedName == CRITERIA_CLASS_FQN } +fun PsiMethodCallExpression.isSuitableForFieldAutoCompleteInAggregation( + parsers: List +): Boolean { + val method = fuzzyResolveMethod() ?: return false + return parsers.any { it.isSuitableForFieldAutoComplete(this, method) } +} + private fun PsiMethodCallExpression.findSpringMongoDbExpression(): PsiMethodCallExpression? { val method = fuzzyResolveMethod() ?: return null if (INTERFACES_WITH_QUERY_METHODS.any { diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParser.kt index f81b1e36..86c33366 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/MatchStageParser.kt @@ -4,6 +4,7 @@ import com.intellij.psi.PsiElement import com.intellij.psi.PsiMethod import com.intellij.psi.PsiMethodCallExpression import com.mongodb.jbplugin.dialects.springcriteria.AGGREGATE_FQN +import com.mongodb.jbplugin.dialects.springcriteria.isCriteriaExpression import com.mongodb.jbplugin.mql.Node import com.mongodb.jbplugin.mql.components.HasFilter import com.mongodb.jbplugin.mql.components.Name @@ -23,6 +24,13 @@ class MatchStageParser( ) ) + override fun isSuitableForFieldAutoComplete( + methodCall: PsiMethodCallExpression, + method: PsiMethod + ): Boolean { + return methodCall.isCriteriaExpression() + } + override fun canParse(stageCallMethod: PsiMethod): Boolean { return stageCallMethod.containingClass?.qualifiedName == AGGREGATE_FQN && stageCallMethod.name == "match" diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt index 7d913641..81e56478 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt @@ -175,6 +175,23 @@ class ProjectStageParser : StageParser { return parseMethodCallWithStringVarArgsAndFields(methodCall, Name.INCLUDE) } + override fun isSuitableForFieldAutoComplete( + methodCall: PsiMethodCallExpression, + method: PsiMethod + ): Boolean { + val methodFqn = method.containingClass?.qualifiedName + // Autocomplete for Aggregation.project("") + return methodFqn == AGGREGATE_FQN && + method.name == "project" || + // Autocomplete for Aggregation.project().andInclude("") + methodFqn == PROJECTION_OPERATION_FQN && + (method.name == "andInclude" || method.name == "andExclude") || + // Autocomplete for Aggregation.project(Fields.fields("")) or, + // Aggregation.project(Fields.from(Fields.field(""))) + methodFqn == FIELDS_FQN && + (method.name == "fields" || method.name == "field") + } + override fun canParse(stageCallMethod: PsiMethod): Boolean { val methodFqn = stageCallMethod.containingClass?.qualifiedName ?: return false // Project stage call might contain chained operations which might result diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/StageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/StageParser.kt index e0f504cc..b349a31d 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/StageParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/StageParser.kt @@ -6,6 +6,10 @@ import com.intellij.psi.PsiMethodCallExpression import com.mongodb.jbplugin.mql.Node interface StageParser { + fun isSuitableForFieldAutoComplete( + methodCall: PsiMethodCallExpression, + method: PsiMethod + ): Boolean fun canParse(stageCallMethod: PsiMethod): Boolean fun parse(stageCall: PsiMethodCallExpression): Node } From e59d87b1d73e502e832d0abd1968c71352f44aad Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 16 Jan 2025 16:20:07 +0100 Subject: [PATCH 8/9] chore: add CHANGELOG entry --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3cffd565..b2bca3b1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ MongoDB plugin for IntelliJ IDEA. ## [Unreleased] ### Added +* [INTELLIJ-173](https://jira.mongodb.org/browse/INTELLIJ-173) Add support for parsing, inspecting and autocompleting in a project stage written using `Aggregation.match` and chained `ProjectionOperations` using `andInclude` and `andExclude`. * [INTELLIJ-172](https://jira.mongodb.org/browse/INTELLIJ-172) Add support for parsing, inspecting and autocompleting in an aggregation written using Spring Data MongoDB (`MongoTemplate.aggregate`, `MongoTemplate.aggregateStream`) and a match stage written using `Aggregation.match`. * [INTELLIJ-179](https://jira.mongodb.org/browse/INTELLIJ-179) Telemetry when Create Index intention is clicked. It can be disabled in the Plugin settings. From f6f960326c653346ca311f4b91d8a01f3d2ec3be Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 16 Jan 2025 16:35:20 +0100 Subject: [PATCH 9/9] chore: minor refactors --- .../ProjectStageParser.kt | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt index 81e56478..224a00c1 100644 --- a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/ProjectStageParser.kt @@ -143,7 +143,7 @@ class ProjectStageParser : StageParser { ): List> { val method = methodCall.fuzzyResolveMethod() ?: return emptyList() return if (method.isVarArgs) { - methodCall.argumentList.expressions.mapNotNull { + methodCall.argumentList.expressions.map { createProjectedFieldNode(fieldExpression = it, projectionName = projectionName) } } else { @@ -154,14 +154,14 @@ class ProjectStageParser : StageParser { } private fun parseAndIncludeCall(methodCall: PsiMethodCallExpression): List> { - // andInclude have two overloads, one which takes string varargs and another + // andInclude has two overloads, one which takes string varargs and another // which takes a Fields object return parseMethodCallWithStringVarArgsAndFields(methodCall, Name.INCLUDE) } private fun parseAndExcludeCall(methodCall: PsiMethodCallExpression): List> { // andExclude only accepts a varargs of strings as arguments - return methodCall.argumentList.expressions.mapNotNull { + return methodCall.argumentList.expressions.map { createProjectedFieldNode(fieldExpression = it, projectionName = Name.EXCLUDE) } } @@ -181,15 +181,21 @@ class ProjectStageParser : StageParser { ): Boolean { val methodFqn = method.containingClass?.qualifiedName // Autocomplete for Aggregation.project("") - return methodFqn == AGGREGATE_FQN && - method.name == "project" || + return ( + methodFqn == AGGREGATE_FQN && + method.name == "project" + ) || // Autocomplete for Aggregation.project().andInclude("") - methodFqn == PROJECTION_OPERATION_FQN && - (method.name == "andInclude" || method.name == "andExclude") || + ( + methodFqn == PROJECTION_OPERATION_FQN && + (method.name == "andInclude" || method.name == "andExclude") + ) || // Autocomplete for Aggregation.project(Fields.fields("")) or, // Aggregation.project(Fields.from(Fields.field(""))) - methodFqn == FIELDS_FQN && - (method.name == "fields" || method.name == "field") + ( + methodFqn == FIELDS_FQN && + (method.name == "fields" || method.name == "field") + ) } override fun canParse(stageCallMethod: PsiMethod): Boolean {