From 6cfdaa4f993743767365a81fb40251214129152d Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 22 Jan 2025 16:01:06 +0100 Subject: [PATCH] feat: add support for parsing, inspecting and autocompleting in a addFields stage written using Spring Data MongoDB INTELLIJ-177 (#125) --- CHANGELOG.md | 1 + ...SpringCriteriaCompletionContributorTest.kt | 302 +++++++ ...iaMongoDbAutocompletionPopupHandlerTest.kt | 307 +++++++- ...gCriteriaFieldCheckLinterInspectionTest.kt | 24 +- .../SpringCriteriaRepository.java | 6 +- .../SpringCriteriaDialectParser.kt | 4 +- .../AddFieldsStageParser.kt | 334 ++++++++ .../ProjectStageParser.kt | 49 +- .../springcriteria/IntegrationTest.kt | 28 +- .../AddFieldsStageParserTest.kt | 742 ++++++++++++++++++ 10 files changed, 1770 insertions(+), 27 deletions(-) create mode 100644 packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/AddFieldsStageParser.kt create mode 100644 packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/AddFieldsStageParserTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b357d7a..ea701fb0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ MongoDB plugin for IntelliJ IDEA. ## [Unreleased] ### Added +* [INTELLIJ-177](https://jira.mongodb.org/browse/INTELLIJ-177) Add support for parsing, inspecting and autocompleting in a addFields stage written using `Aggregation.addFields` and chained `AddFieldsOperation`s using `addFieldWithValue`, `addFieldWithValueOf`, `addField().withValue()` and `addField().withValueOf()`. Parsing boxed Java values is not supported yet. * [INTELLIJ-174](https://jira.mongodb.org/browse/INTELLIJ-174) Add support for parsing, inspecting and autocompleting in a sort stage written using `Aggregation.sort` and chained `SortOperation`s using `and`. All the overloads of creating a `Sort` object are supported. * [INTELLIJ-188](https://jira.mongodb.org/browse/INTELLIJ-188) Support for generating sort in the query generator. * [INTELLIJ-186](https://jira.mongodb.org/browse/INTELLIJ-186) Support for parsing Sorts in the Java Driver. 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 ba83b14d..f33f740a 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 @@ -1531,4 +1531,306 @@ class Repository { }, ) } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.domain.Sort; +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.addFields().addFieldWithValueOf("addedField", "") + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in Aggregation#addFields#addFieldsWithValueOf 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.domain.Sort; +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.addFields().addFieldWithValueOf("addedField", Fields.field("")) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in Fields#field call passed to Aggregation#addFields#addFieldsWithValueOf 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.domain.Sort; +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.addFields().addField("addedField").withValueOf("") + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in Aggregation#addFields#withValueOf 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.domain.Sort; +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.addFields().addField("addedField").withValueOf(Fields.field("")) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in Fields#field call passed to Aggregation#addFields#withValueOf 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 fbd3e1ae..6b04d65b 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 @@ -14,7 +14,6 @@ import org.junit.jupiter.api.Assertions.assertTrue import org.mockito.Mockito.`when` import org.mockito.kotlin.eq -@Suppress("TOO_LONG_FUNCTION") @CodeInsightTest class SpringCriteriaMongoDbAutocompletionPopupHandlerTest { @ParsingTest( @@ -1478,4 +1477,310 @@ class Repository { }, ) } + + @ParsingTest( + fileName = "Repository.java", + value = """ +import org.springframework.data.domain.Sort; +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.addFields().addFieldWithValueOf("addedField", ) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in Aggregation#addFields#addFieldsWithValueOf 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.domain.Sort; +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.addFields().addFieldWithValueOf("addedField", Fields.field()) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in Fields#field call passed to Aggregation#addFields#addFieldsWithValueOf 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.domain.Sort; +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.addFields().addField("addedField").withValueOf() + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in Aggregation#addFields#withValueOf 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.domain.Sort; +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.addFields().addField("addedField").withValueOf(Fields.field()) + ), + Book.class, + Book.class + ).getMappedResults(); + } +} + """, + ) + fun `should autocomplete fields from the current namespace in Fields#field call passed to Aggregation#addFields#withValueOf 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/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 488ee474..da55f8d5 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 @@ -169,7 +169,7 @@ class BookRepository { import org.springframework.data.domain.Sort; 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.aggregation.Fields;import org.springframework.data.mongodb.core.mapping.Document; import java.util.List; import static org.springframework.data.mongodb.core.query.Query.query; @@ -270,7 +270,27 @@ class BookRepository { ) ).descending() ) - ) + ), + Aggregation.addFields() + // no inspection here as no field reference + .addFieldWithValue("addedField", "released") + // field reference as computed value + .addFieldWithValueOf("addedField", "released") + .addFieldWithValueOf("addedField", releasedAsVariable) + .addFieldWithValueOf("addedField", releasedFromMethodCall()) + .addFieldWithValueOf("addedField", Fields.field("released")) + .addFieldWithValueOf("addedField", Fields.field(releasedAsVariable)) + .addFieldWithValueOf("addedField", Fields.field(releasedFromMethodCall())) + // no inspection here as no field reference + .addField("addedField").withValue("released") + // field reference as computed value + .addField("addedField").withValueOf("released") + .addField("addedField").withValueOf(releasedAsVariable) + .addField("addedField").withValueOf(releasedFromMethodCall()) + .addField("addedField").withValueOf(Fields.field("released")) + .addField("addedField").withValueOf(Fields.field(releasedAsVariable)) + .addField("addedField").withValueOf(Fields.field(releasedFromMethodCall())) + .build() ), Book.class, Book.class 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 191eef17..10877210 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 @@ -41,7 +41,11 @@ private List allMoviesWithRatingAtLeastAgg(int rating) { Aggregation.sort(Sort.by("rated", "qwe")), Aggregation.sort(Sort.by(Sort.Direction.ASC, "rates")), Aggregation.sort(Sort.by(Sort.Order.by("rated"), Sort.Order.by("ratedd"))), - Aggregation.sort(Sort.by(List.of(Sort.Order.by("rated"), Sort.Order.by("rateds")))) + Aggregation.sort(Sort.by(List.of(Sort.Order.by("rated"), Sort.Order.by("rateds")))), + Aggregation.addFields().addFieldWithValueOf("addedField", "value").build(), + Aggregation.addFields().addFieldWithValueOf("addedField", Fields.field("qwe")).build(), + Aggregation.addFields().addField("addedField").withValueOf("rateds").build(), + Aggregation.addFields().addField("addedField").withValueOf(Fields.field("asd")).build() ), Movie.class, Movie.class 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 915e6880..5c412f88 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.aggregationstageparsers.AddFieldsStageParser import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.MatchStageParser import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.ProjectStageParser import com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers.SortStageParser @@ -33,7 +34,8 @@ object SpringCriteriaDialectParser : DialectParser { MatchStageParser(::parseFilterRecursively), ProjectStageParser(), UnwindStageParser(), - SortStageParser() + SortStageParser(), + AddFieldsStageParser(), ) override fun isCandidateForQuery(source: PsiElement) = diff --git a/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/AddFieldsStageParser.kt b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/AddFieldsStageParser.kt new file mode 100644 index 00000000..3d0f3d1c --- /dev/null +++ b/packages/mongodb-dialects/spring-criteria/src/main/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/AddFieldsStageParser.kt @@ -0,0 +1,334 @@ +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.collectTypeUntil +import com.mongodb.jbplugin.dialects.javadriver.glossary.fuzzyResolveMethod +import com.mongodb.jbplugin.dialects.javadriver.glossary.resolveToMethodCallExpression +import com.mongodb.jbplugin.dialects.javadriver.glossary.tryToResolveAsConstant +import com.mongodb.jbplugin.dialects.javadriver.glossary.tryToResolveAsConstantString +import com.mongodb.jbplugin.dialects.springcriteria.AGGREGATE_FQN +import com.mongodb.jbplugin.mql.BsonAny +import com.mongodb.jbplugin.mql.ComputedBsonType +import com.mongodb.jbplugin.mql.Node +import com.mongodb.jbplugin.mql.components.HasAddedFields +import com.mongodb.jbplugin.mql.components.HasFieldReference +import com.mongodb.jbplugin.mql.components.HasFieldReference.FieldReference +import com.mongodb.jbplugin.mql.components.HasValueReference +import com.mongodb.jbplugin.mql.components.HasValueReference.ValueReference +import com.mongodb.jbplugin.mql.components.Name +import com.mongodb.jbplugin.mql.components.Named +import com.mongodb.jbplugin.mql.toBsonType + +internal const val ADD_FIELDS_OPERATION_FQN = "org.springframework.data.mongodb.core.aggregation.AddFieldsOperation" +internal const val ADD_FIELDS_OPERATION_BUILDER_FQN = "$ADD_FIELDS_OPERATION_FQN.AddFieldsOperationBuilder" +internal const val VALUE_APPENDER_FQN = "$ADD_FIELDS_OPERATION_BUILDER_FQN.ValueAppender" + +class AddFieldsStageParser : StageParser { + override fun isSuitableForFieldAutoComplete( + methodCall: PsiMethodCallExpression, + method: PsiMethod + ): Boolean { + return methodCall.isAddFieldWithValueOfCall() || + methodCall.isWithValueOfCall() || + methodCall.isFieldCreationCallInsideAddFieldsChain() + } + + override fun canParse(stageCallMethod: PsiMethod): Boolean { + return listOf( + AGGREGATE_FQN, + ADD_FIELDS_OPERATION_FQN, + ADD_FIELDS_OPERATION_BUILDER_FQN, + VALUE_APPENDER_FQN + ).contains(stageCallMethod.containingClass?.qualifiedName) && + listOf( + "addFields", + "addFieldWithValue", + "addFieldWithValueOf", + "addField", + "withValue", + "withValueOf", + "build" + ).contains(stageCallMethod.name) + } + + override fun parse(stageCall: PsiMethodCallExpression): Node { + val allChainedCalls = stageCall.gatherChainedCalls() + return createAddFieldsNode( + source = stageCall, + addedFields = allChainedCalls.mapIndexedNotNull { index, methodCall -> + val method = methodCall.fuzzyResolveMethod() ?: return@mapIndexedNotNull null + when (method.name) { + // Nothing to parse in these method calls + "addFields", "build" -> null + // addField contains the field name that is added but that is parsed alongside + // the parsing of withValue and withValueOf which is why we ignore it here + "addField" -> null + "addFieldWithValue" -> parseAddFieldWithValue(methodCall) + "addFieldWithValueOf" -> parseAddFieldWithValueOf(methodCall) + "withValue" -> parseWithValue( + // withValue can only be chained on a ValueAppender instance and thus addField + // is expected to be right after the withValue call in the chain + addFieldCall = allChainedCalls.getOrNull(index + 1)?.takeIf { + it.isAddFieldCall() + }, + withValueCall = methodCall + ) + "withValueOf" -> parseWithValueOf( + // withValueOf can only be chained on a ValueAppender instance and thus addField + // is expected to be right after the withValue call in the chain + addFieldCall = allChainedCalls.getOrNull(index + 1)?.takeIf { + it.isAddFieldCall() + }, + withValueCall = methodCall + ) + else -> null + } + } + ) + } + + /** + * Parses a PsiMethodCallExpression representing withValue() method. + * withValue() itself accepts a value which can either be a Java primitive type or a reference + * type and is chained on a ValueAppender instance built using addField() method which accepts + * a field as a String. To parse withValue() it is imperative to also parse addField() call. + * + * Note: Current we support parsing values which are Java primitives and can be resolved during + * build time. + */ + private fun parseWithValue( + addFieldCall: PsiMethodCallExpression?, + withValueCall: PsiMethodCallExpression + ): Node { + val fieldReference = parseFieldExpression( + addFieldCall?.argumentList?.expressions?.getOrNull(0) + ) + val valueReference = parseValueExpressionAsNormalValue( + withValueCall.argumentList.expressions.getOrNull(0) + ) + + return createAddFieldNode( + source = withValueCall, + fieldReference = fieldReference, + valueReference = valueReference + ) + } + + /** + * Parses a PsiMethodCallExpression representing withValueOf() method. + * withValueOf() accepts a value which could point to: + * 1. existing field in the database + * 2. a MongoDB expression that evaluates to some value + * Note: Currently we support parsing values which points to a field either directly as a String + * or using a `Field object + * + * withValueOf() is chained on a ValueAppender instance built using addField() method which + * accepts a field as a String. To parse withValue() it is imperative to also parse addField(). + */ + private fun parseWithValueOf( + addFieldCall: PsiMethodCallExpression?, + withValueCall: PsiMethodCallExpression + ): Node { + val fieldReference = parseFieldExpression( + addFieldCall?.argumentList?.expressions?.getOrNull(0) + ) + val valueReference = parseValueExpressionAsField( + withValueCall.argumentList.expressions.getOrNull(0) + ) + + return createAddFieldNode( + source = withValueCall, + fieldReference = fieldReference, + valueReference = valueReference + ) + } + + /** + * Parses a PsiMethodCallExpression representing addFieldWithValue() method. + * addFieldWithValue() accepts a field as a String and a value which can either be a Java + * primitive type or a reference type. + * Note: Current we support parsing values which are Java primitives and can be resolved during + * build time. + */ + private fun parseAddFieldWithValue(methodCall: PsiMethodCallExpression): Node { + val fieldReference = parseFieldExpression( + methodCall.argumentList.expressions.getOrNull(0) + ) + val valueReference = parseValueExpressionAsNormalValue( + methodCall.argumentList.expressions.getOrNull(1) + ) + + return createAddFieldNode( + source = methodCall, + fieldReference = fieldReference, + valueReference = valueReference + ) + } + + /** + * Parses a PsiMethodCallExpression representing addFieldWithValueOf() method. + * addFieldWithValueOf() accepts a field as a String and a value which could point to: + * 1. existing field in the database + * 2. a MongoDB expression that evaluates to some value + * Note: Currently we support parsing values which points to a field either directly as a String + * or using a `Field object + */ + private fun parseAddFieldWithValueOf(methodCall: PsiMethodCallExpression): Node { + val fieldReference = parseFieldExpression( + methodCall.argumentList.expressions.getOrNull(0) + ) + val valueReference = parseValueExpressionAsField( + methodCall.argumentList.expressions.getOrNull(1) + ) + + return createAddFieldNode( + source = methodCall, + fieldReference = fieldReference, + valueReference = valueReference + ) + } + + private fun parseFieldExpression(fieldExpression: PsiExpression?): FieldReference { + val resolvedFieldAddedInCall = fieldExpression?.tryToResolveAsConstantString() + return if (resolvedFieldAddedInCall == null) { + HasFieldReference.Unknown as FieldReference + } else { + HasFieldReference.Computed( + source = fieldExpression, + fieldName = resolvedFieldAddedInCall, + displayName = resolvedFieldAddedInCall, + ) + } + } + + private fun parseValueExpressionAsNormalValue(valueExpression: PsiExpression?): ValueReference { + val (wasResolved, resolvedValue) = valueExpression?.tryToResolveAsConstant() + ?: (false to null) + + if (!wasResolved || valueExpression == null) { + return HasValueReference.Unknown as ValueReference + } + + return HasValueReference.Constant( + source = valueExpression, + value = resolvedValue, + type = resolvedValue?.javaClass.toBsonType(), + ) + } + + private fun parseValueExpressionAsField(valueExpression: PsiExpression?): ValueReference { + if (valueExpression == null) { + return HasValueReference.Unknown as ValueReference + } + + val fieldExpressionReferencedInValueExpression = if (valueExpression.isFieldObject()) { + // Since valueExpression is a `Field` object, the value passed to the object is likely + // a field in mongodb schema so we pluck that out. + valueExpression.resolveFieldStringExpressionFromFieldObject() + } else { + valueExpression + } + val resolvedFieldReferencedInValueExpression = fieldExpressionReferencedInValueExpression + ?.tryToResolveAsConstantString() + + return if (resolvedFieldReferencedInValueExpression == null) { + HasValueReference.Unknown as ValueReference + } else { + HasValueReference.Computed( + source = valueExpression, + type = ComputedBsonType( + baseType = BsonAny, + expression = Node( + source = fieldExpressionReferencedInValueExpression, + components = listOf( + HasFieldReference( + HasFieldReference.FromSchema( + source = fieldExpressionReferencedInValueExpression, + fieldName = resolvedFieldReferencedInValueExpression.trim('$'), + displayName = "\$${resolvedFieldReferencedInValueExpression.trim( + '$' + )}" + ) + ) + ) + ) + ) + ) + } + } + + private fun createAddFieldNode( + source: PsiElement, + fieldReference: FieldReference, + valueReference: ValueReference, + ): Node { + return Node( + source = source, + components = listOf( + Named(Name.ADD_FIELDS), + HasFieldReference(fieldReference), + HasValueReference(valueReference) + ) + ) + } + + private fun createAddFieldsNode( + source: PsiMethodCallExpression, + addedFields: List> + ): Node { + return Node( + source = source, + components = listOf( + Named(Name.ADD_FIELDS), + HasAddedFields(addedFields) + ) + ) + } +} + +fun PsiMethodCallExpression.isAddFieldWithValueOfCall(): Boolean { + val method = fuzzyResolveMethod() ?: return false + return method.name == "addFieldWithValueOf" && + method.containingClass?.qualifiedName == ADD_FIELDS_OPERATION_BUILDER_FQN +} + +fun PsiMethodCallExpression.isWithValueOfCall(): Boolean { + val method = fuzzyResolveMethod() ?: return false + return method.name == "withValueOf" && + method.containingClass?.qualifiedName == VALUE_APPENDER_FQN +} + +/** + * Confirms if a `Field` object creation call was inside one of: + * 1. AddFieldsOperationBuilder.addFieldWithValueOf() + * 2. ValueAppender.withValueOf() + */ +fun PsiMethodCallExpression.isFieldCreationCallInsideAddFieldsChain(): Boolean { + val method = fuzzyResolveMethod() ?: return false + val allParents = method.collectTypeUntil( + PsiMethodCallExpression::class.java, + PsiMethod::class.java + ) + return method.isFieldCreationMethod() && + allParents.any { + isAddFieldWithValueOfCall() || isWithValueOfCall() + } +} + +fun PsiMethodCallExpression.isAddFieldCall(): Boolean { + val method = fuzzyResolveMethod() ?: return false + return method.name == "addField" && + method.containingClass?.qualifiedName == ADD_FIELDS_OPERATION_BUILDER_FQN +} + +/** + * Confirms if a PsiExpression represents a `Field` object + */ +fun PsiExpression.isFieldObject(): Boolean { + return resolveToMethodCallExpression { _, method -> + method.isFieldCreationMethod() + } != null +} 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 d0a39088..6014cd18 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 @@ -73,26 +73,12 @@ class ProjectStageParser : StageParser { 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) - } + val fieldExpression = fieldObjectExpression.resolveFieldStringExpressionFromFieldObject() + ?: return null + return createProjectedFieldNode( + fieldExpression = fieldExpression, + projectionName = projectionName + ) } /** @@ -224,6 +210,29 @@ class ProjectStageParser : StageParser { } } +fun PsiExpression.resolveFieldStringExpressionFromFieldObject(): PsiExpression? { + val resolvedFieldMethodCall = resolveToMethodCallExpression { + _, + fieldObjectMethod + -> + fieldObjectMethod.isFieldCreationMethod() + } ?: 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. + return if (resolvedFieldMethodCall.argumentList.expressions.size == 2) { + null + } else { + resolvedFieldMethodCall.argumentList.expressions.getOrNull(0) + } +} + +fun PsiMethod.isFieldCreationMethod(): Boolean { + return containingClass?.qualifiedName == FIELDS_FQN && name == "field" +} + /** * From a PsiMethodCallExpression, it attempts to travel backwards the chain(assuming there is one) * while gathering other PsiMethodCallExpressions it comes across, until there is no further 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 3cf39b78..411b9044 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 @@ -24,17 +24,16 @@ import com.intellij.psi.util.PsiUtil import com.intellij.psi.util.childrenOfType import com.intellij.testFramework.IdeaTestUtil import com.intellij.testFramework.IndexingTestUtil -import com.intellij.testFramework.LightProjectDescriptor.SetupHandler import com.intellij.testFramework.PsiTestUtil import com.intellij.testFramework.TestApplicationManager import com.intellij.testFramework.TestDataPath import com.intellij.testFramework.fixtures.CodeInsightTestFixture import com.intellij.testFramework.fixtures.DefaultLightProjectDescriptor -import com.intellij.testFramework.fixtures.DefaultLightProjectDescriptor.addJetBrainsAnnotations import com.intellij.testFramework.fixtures.IdeaTestFixtureFactory import com.mongodb.assertions.Assertions.assertNotNull import com.mongodb.jbplugin.mql.Component import com.mongodb.jbplugin.mql.Node +import com.mongodb.jbplugin.mql.components.HasAddedFields import com.mongodb.jbplugin.mql.components.HasAggregation import com.mongodb.jbplugin.mql.components.HasCollectionReference import com.mongodb.jbplugin.mql.components.HasFieldReference @@ -403,6 +402,31 @@ fun Node.sortN( sort.assertions() } +fun Node.addedFieldN( + n: Int, + name: Name? = null, + stageIndex: Int? = null, + assertions: Node.() -> Unit = { + }, +) { + val addedFields = component>() + assertNotNull(addedFields) + + val addedField = addedFields!!.children[n] + + if (name != null) { + val qname = addedField.component() + assertNotEquals(null, qname) { + "StageIndex: $stageIndex, AddedFieldIndex: $n :: Expected a named operation with name $name but null found." + } + assertEquals(name, qname?.name) { + "StageIndex: $stageIndex, AddedFieldIndex: $n :: Expected a named operation with name $name but $qname found." + } + } + + addedField.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/AddFieldsStageParserTest.kt b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/AddFieldsStageParserTest.kt new file mode 100644 index 00000000..73f20bff --- /dev/null +++ b/packages/mongodb-dialects/spring-criteria/src/test/kotlin/com/mongodb/jbplugin/dialects/springcriteria/aggregationstageparsers/AddFieldsStageParserTest.kt @@ -0,0 +1,742 @@ +package com.mongodb.jbplugin.dialects.springcriteria.aggregationstageparsers + +import com.intellij.psi.PsiElement +import com.intellij.psi.PsiExpression +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.addedFieldN +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.stageN +import com.mongodb.jbplugin.dialects.springcriteria.value +import com.mongodb.jbplugin.mql.BsonAnyOf +import com.mongodb.jbplugin.mql.BsonBoolean +import com.mongodb.jbplugin.mql.BsonDouble +import com.mongodb.jbplugin.mql.BsonInt32 +import com.mongodb.jbplugin.mql.BsonInt64 +import com.mongodb.jbplugin.mql.BsonNull +import com.mongodb.jbplugin.mql.BsonString +import com.mongodb.jbplugin.mql.BsonType +import com.mongodb.jbplugin.mql.components.HasAddedFields +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.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 +import org.junit.jupiter.api.Assertions.assertTrue + +@IntegrationTest +class AddFieldsStageParserTest { + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.AddFieldsOperation; +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; + } + + AddFieldsOperationBuilder getPartialAddFields() { + return Aggregation.addFields(); + } + + AddFieldsOperation getCompleteAddFields() { + return Aggregation.addFields().build(); + } + + AddFieldsOperation getChainedAddFields() { + return Aggregation.addFields().addField().withValue().build(); + } + + public AggregationResults allReleasedBooks() { + var partialAddFieldsAsVariable = Aggregation.addFields(); + AddFieldsOperation completeAddFieldsAsVariable = Aggregation.addFields().build(); + AddFieldsOperation chainedAddFieldsAsVariable = Aggregation.addFields().addFieldWithValueOf().build(); + return template.aggregate( + Aggregation.newAggregation( + // this call is not correct but our parser is able to parse it + Aggregation.addFields(), + partialAddFieldsAsVariable, + getPartialAddFields(), + // A complete call as it terminates with build + Aggregation.addFields().build(), + completeAddFieldsAsVariable, + getCompleteAddFields(), + // Chained builder calls but empty + Aggregation.addFields().addFieldWithValue().build(), + chainedAddFieldsAsVariable, + getChainedAddFields() + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse an empty addFields stage`(psiFile: PsiFile) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + parseAndAssertForAddFieldsStagesWithConstantValue(query, 9, emptyList()) + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.AddFieldsOperation; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; +import java.util.List; +import java.util.Map; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + Strin getField() { + return "field0"; + } + + Strin getStringValue() { + return "value0"; + } + + int getIntValue() { + return 1; + } + + long getLongValue() { + return 1L; + } + + double getDoubleValue() { + return 1.0; + } + + boolean getBooleanValue() { + return true; + } + + AddFieldsOperation getAddFieldsOperation() { + return Aggregation.addFields().addFieldWithValue(getField(), getStringValue()).build(); + } + + AddFieldsOperation getAddFieldsChain() { + int intAsVariable = 1; + long longAsVariable = 1L; + double doubleAsVariable = 1.0; + boolean booleanAsVariable = true; + return Aggregation.addFields() + .addFieldWithValue("field0", 1) + .addFieldWithValue("field0", intAsVariable) + .addFieldWithValue("field0", getIntValue()) + .addFieldWithValue("field0", 1L) + .addFieldWithValue("field0", longAsVariable) + .addFieldWithValue("field0", getLongValue()) + .addFieldWithValue("field0", 1.0) + .addFieldWithValue("field0", doubleAsVariable) + .addFieldWithValue("field0", getDoubleValue()) + .addFieldWithValue("field0", true) + .addFieldWithValue("field0", booleanAsVariable) + .addFieldWithValue("field0", getBooleanValue()) + .build(); + } + + public AggregationResults allReleasedBooks() { + String fieldAsVariable = "field0"; + String stringValueAsVariable = "value0"; + int intAsVariable = 1; + long longAsVariable = 1L; + double doubleAsVariable = 1.0; + boolean booleanAsVariable = true; + + AddFieldsOperation addFieldsOperationAsVariable = Aggregation.addFields() + .addFieldWithValue(fieldAsVariable, stringValueAsVariable) + .build(); + AddFieldsOperation addFieldsChainAsVariable = Aggregation.addFields() + .addFieldWithValue("field0", 1) + .addFieldWithValue("field0", intAsVariable) + .addFieldWithValue("field0", getIntValue()) + .addFieldWithValue("field0", 1L) + .addFieldWithValue("field0", longAsVariable) + .addFieldWithValue("field0", getLongValue()) + .addFieldWithValue("field0", 1.0) + .addFieldWithValue("field0", doubleAsVariable) + .addFieldWithValue("field0", getDoubleValue()) + .addFieldWithValue("field0", true) + .addFieldWithValue("field0", booleanAsVariable) + .addFieldWithValue("field0", getBooleanValue()) + .build(); + + return template.aggregate( + Aggregation.newAggregation( + // addFieldWithValue variants + Aggregation.addFields().addFieldWithValue("field0", "value0"), // incorrect because no build but our parser should be able to parse this + Aggregation.addFields().addFieldWithValue("field0", "value0").build(), + Aggregation.addFields().addFieldWithValue(fieldAsVariable, stringValueAsVariable).build(), + Aggregation.addFields().addFieldWithValue(getField(), getStringValue()).build(), + addFieldsOperationAsVariable, + getAddFieldsOperation(), + // chained addFieldWithValue with build time constant values other than string + Aggregation.addFields() + .addFieldWithValue("field0", 1) + .addFieldWithValue("field0", intAsVariable) + .addFieldWithValue("field0", getIntValue()) + .addFieldWithValue("field0", 1L) + .addFieldWithValue("field0", longAsVariable) + .addFieldWithValue("field0", getLongValue()) + .addFieldWithValue("field0", 1.0) + .addFieldWithValue("field0", doubleAsVariable) + .addFieldWithValue("field0", getDoubleValue()) + .addFieldWithValue("field0", true) + .addFieldWithValue("field0", booleanAsVariable) + .addFieldWithValue("field0", getBooleanValue()) + .build(), + addFieldsChainAsVariable, + getAddFieldsChain() + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse an addFields stage with added fields via addFieldWithValue`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + val stringFieldValueExpectation = + 1 to listOf(Triple("field0", "value0", BsonAnyOf(BsonString, BsonNull))) + val intFieldValueExpectation = Triple("field0", 1, BsonAnyOf(BsonInt32, BsonNull)) + val longFieldValueExpectation = Triple("field0", 1L, BsonAnyOf(BsonInt64, BsonNull)) + val doubleFieldValueExpectation = Triple("field0", 1.0, BsonAnyOf(BsonDouble, BsonNull)) + val booleanFieldValueExpectation = Triple("field0", true, BsonAnyOf(BsonBoolean, BsonNull)) + val chainedAddFieldExpectation = 12 to listOf( + booleanFieldValueExpectation, + booleanFieldValueExpectation, + booleanFieldValueExpectation, + doubleFieldValueExpectation, + doubleFieldValueExpectation, + doubleFieldValueExpectation, + longFieldValueExpectation, + longFieldValueExpectation, + longFieldValueExpectation, + intFieldValueExpectation, + intFieldValueExpectation, + intFieldValueExpectation, + ) + parseAndAssertForAddFieldsStagesWithConstantValue( + query, + 9, + listOf( + stringFieldValueExpectation, + stringFieldValueExpectation, + stringFieldValueExpectation, + stringFieldValueExpectation, + stringFieldValueExpectation, + stringFieldValueExpectation, + chainedAddFieldExpectation, + chainedAddFieldExpectation, + chainedAddFieldExpectation, + ) + ) + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.AddFieldsOperation; +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; +import java.util.Map; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + String getFieldValueAsString() { + return "field0FromSchema"; + } + + Field getFieldValueAsField() { + return Fields.field("field0FromSchema"); + } + + AddFieldsOperation getSingleOperation() { + return Aggregation.addFields().addFieldWithValueOf("field0", getFieldValueAsField()).build(); + } + + public AggregationResults allReleasedBooks() { + String fieldAsString = "field0FromSchema"; + Field fieldAsField = Fields.field("field0FromSchema"); + AddFieldsOperation singleOperation = Aggregation.addFields().addFieldWithValueOf("field0", "field0FromSchema").build(); + return template.aggregate( + Aggregation.newAggregation( + Aggregation.addFields().addFieldWithValueOf("field0", "field0FromSchema"), // incorrect but we should be able to parse this + singleOperation, + getSingleOperation(), + Aggregation.addFields() + .addFieldWithValueOf("field0", fieldAsString) + .addFieldWithValueOf("field1", getFieldValueAsString()) + .addFieldWithValueOf("field2", Fields.field("field0FromSchema")) + .addFieldWithValueOf("field3", fieldAsField) + .addFieldWithValueOf("field4", getFieldValueAsField()) + .build() + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse an addFields stage with added fields via addFieldWithValueOf`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + parseAndAssertForAddFieldsStagesWithComputedValue( + query, + 4, + listOf( + 1 to listOf("field0" to "field0FromSchema"), + 1 to listOf("field0" to "field0FromSchema"), + 1 to listOf("field0" to "field0FromSchema"), + 5 to listOf( + "field4" to "field0FromSchema", + "field3" to "field0FromSchema", + "field2" to "field0FromSchema", + "field1" to "field0FromSchema", + "field0" to "field0FromSchema", + ) + ), + ) + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.AddFieldsOperation; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.mapping.Document; +import java.util.List; +import java.util.Map; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + Strin getField() { + return "field0"; + } + + Strin getStringValue() { + return "value0"; + } + + int getIntValue() { + return 1; + } + + long getLongValue() { + return 1L; + } + + double getDoubleValue() { + return 1.0; + } + + boolean getBooleanValue() { + return true; + } + + AddFieldsOperation getAddFieldsOperation() { + return Aggregation.addFields().addField(getField()).withValue(getStringValue()).build(); + } + + AddFieldsOperation getAddFieldsChain() { + int intAsVariable = 1; + long longAsVariable = 1L; + double doubleAsVariable = 1.0; + boolean booleanAsVariable = true; + return Aggregation.addFields() + .addField("field0").withValue(1) + .addField("field0").withValue(intAsVariable) + .addField("field0").withValue(getIntValue()) + .addField("field0").withValue(1L) + .addField("field0").withValue(longAsVariable) + .addField("field0").withValue(getLongValue()) + .addField("field0").withValue(1.0) + .addField("field0").withValue(doubleAsVariable) + .addField("field0").withValue(getDoubleValue()) + .addField("field0").withValue(true) + .addField("field0").withValue(booleanAsVariable) + .addField("field0").withValue(getBooleanValue()) + .build(); + } + + public AggregationResults allReleasedBooks() { + String fieldAsVariable = "field0"; + String stringValueAsVariable = "value0"; + int intAsVariable = 1; + long longAsVariable = 1L; + double doubleAsVariable = 1.0; + boolean booleanAsVariable = true; + + AddFieldsOperation addFieldsOperationAsVariable = Aggregation.addFields() + .addField(fieldAsVariable).withValue(stringValueAsVariable) + .build(); + + AddFieldsOperation addFieldsChainAsVariable = Aggregation.addFields() + .addField("field0").withValue(1) + .addField("field0").withValue(intAsVariable) + .addField("field0").withValue(getIntValue()) + .addField("field0").withValue(1L) + .addField("field0").withValue(longAsVariable) + .addField("field0").withValue(getLongValue()) + .addField("field0").withValue(1.0) + .addField("field0").withValue(doubleAsVariable) + .addField("field0").withValue(getDoubleValue()) + .addField("field0").withValue(true) + .addField("field0").withValue(booleanAsVariable) + .addField("field0").withValue(getBooleanValue()) + .build(); + + return template.aggregate( + Aggregation.newAggregation( + // addFieldWithValue variants + Aggregation.addFields().addField("field0").withValue("value0"), // incorrect because no build but our parser should be able to parse this + Aggregation.addFields().addField("field0").withValue("value0").build(), + Aggregation.addFields().addField(fieldAsVariable).withValue(stringValueAsVariable).build(), + Aggregation.addFields().addField(getField()).withValue(getStringValue()).build(), + addFieldsOperationAsVariable, + getAddFieldsOperation(), + // chained addFieldWithValue with build time constant values other than string + Aggregation.addFields() + .addField("field0").withValue(1) + .addField("field0").withValue(intAsVariable) + .addField("field0").withValue(getIntValue()) + .addField("field0").withValue(1L) + .addField("field0").withValue(longAsVariable) + .addField("field0").withValue(getLongValue()) + .addField("field0").withValue(1.0) + .addField("field0").withValue(doubleAsVariable) + .addField("field0").withValue(getDoubleValue()) + .addField("field0").withValue(true) + .addField("field0").withValue(booleanAsVariable) + .addField("field0").withValue(getBooleanValue()) + .build(), + addFieldsChainAsVariable, + getAddFieldsChain() + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse an addFields stage with added fields via addField#withValue`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + val stringFieldValueExpectation = + 1 to listOf(Triple("field0", "value0", BsonAnyOf(BsonString, BsonNull))) + val intFieldValueExpectation = Triple("field0", 1, BsonAnyOf(BsonInt32, BsonNull)) + val longFieldValueExpectation = Triple("field0", 1L, BsonAnyOf(BsonInt64, BsonNull)) + val doubleFieldValueExpectation = Triple("field0", 1.0, BsonAnyOf(BsonDouble, BsonNull)) + val booleanFieldValueExpectation = Triple("field0", true, BsonAnyOf(BsonBoolean, BsonNull)) + val chainedAddFieldExpectation = 12 to listOf( + booleanFieldValueExpectation, + booleanFieldValueExpectation, + booleanFieldValueExpectation, + doubleFieldValueExpectation, + doubleFieldValueExpectation, + doubleFieldValueExpectation, + longFieldValueExpectation, + longFieldValueExpectation, + longFieldValueExpectation, + intFieldValueExpectation, + intFieldValueExpectation, + intFieldValueExpectation, + ) + parseAndAssertForAddFieldsStagesWithConstantValue( + query, + 9, + listOf( + stringFieldValueExpectation, + stringFieldValueExpectation, + stringFieldValueExpectation, + stringFieldValueExpectation, + stringFieldValueExpectation, + stringFieldValueExpectation, + chainedAddFieldExpectation, + chainedAddFieldExpectation, + chainedAddFieldExpectation, + ) + ) + } + + @ParsingTest( + fileName = "Book.java", + """ +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.AddFieldsOperation; +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; +import java.util.Map; + +@Document +record Book() {} + +class Repository { + private final MongoTemplate template; + + public Repository(MongoTemplate template) { + this.template = template; + } + + String getFieldValueAsString() { + return "field0FromSchema"; + } + + Field getFieldValueAsField() { + return Fields.field("field0FromSchema"); + } + + AddFieldsOperation getSingleOperation() { + return Aggregation.addFields().addField("field0").withValueOf(getFieldValueAsField()).build(); + } + + public AggregationResults allReleasedBooks() { + String field = "field0"; + String fieldAsString = "field0FromSchema"; + Field fieldAsField = Fields.field("field0FromSchema"); + AddFieldsOperation singleOperation = Aggregation.addFields().addField(field).withValueOf(fieldAsString).build(); + return template.aggregate( + Aggregation.newAggregation( + Aggregation.addFields().addField("field0").withValueOf("field0FromSchema"), // incorrect but we should be able to parse this + singleOperation, + getSingleOperation(), + Aggregation.addFields() + .addField("field0").withValueOf(fieldAsString) + .addField("field1").withValueOf(getFieldValueAsString()) + .addField("field2").withValueOf(Fields.field("field0FromSchema")) + .addField("field3").withValueOf(fieldAsField) + .addField("field4").withValueOf(getFieldValueAsField()) + .build() + ), + Book.class, + Book.class + ); + } +} + """ + ) + fun `should be able to parse an addFields stage with added fields via addField#withValueOf`( + psiFile: PsiFile + ) { + val query = psiFile.getQueryAtMethod("Repository", "allReleasedBooks") + parseAndAssertForAddFieldsStagesWithComputedValue( + query, + 4, + listOf( + 1 to listOf("field0" to "field0FromSchema"), + 1 to listOf("field0" to "field0FromSchema"), + 1 to listOf("field0" to "field0FromSchema"), + 5 to listOf( + "field4" to "field0FromSchema", + "field3" to "field0FromSchema", + "field2" to "field0FromSchema", + "field1" to "field0FromSchema", + "field0" to "field0FromSchema", + ) + ), + ) + } + + companion object { + fun parseAndAssertForAddFieldsStagesWithConstantValue( + query: PsiExpression, + expectedAddFieldsStagesCounts: Int, + addFieldsStagesExpectation: List>>> + ) { + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + component> { + assertEquals( + expectedAddFieldsStagesCounts, + children.size, + "Expected $expectedAddFieldsStagesCounts aggregation stages but found ${children.size}" + ) + } + + addFieldsStagesExpectation.forEachIndexed { + stageIndex, + (expectedAddedFieldsCount, addedFieldsExpectations) + -> + stageN(stageIndex, Name.ADD_FIELDS) { + component> { + assertEquals( + expectedAddedFieldsCount, + children.size, + "StageIndex $stageIndex :: Expected $expectedAddedFieldsCount added fields but found ${children.size}" + ) + } + + addedFieldsExpectations.forEachIndexed { + index, + (expectedFieldName, expectedValue, expectedType) + -> + addedFieldN(index, Name.ADD_FIELDS, stageIndex) { + field> { + assertEquals( + expectedFieldName, + fieldName, + "StageIndex $stageIndex, AddedFieldIndex: $index :: Expected field with name $expectedFieldName but found field with name $fieldName" + ) + } + value> { + assertEquals( + expectedType, + type, + "StageIndex $stageIndex, AddedFieldIndex: $index :: Expected value type $expectedType but found $type" + ) + assertEquals( + expectedValue, + value, + "StageIndex $stageIndex, AddedFieldIndex: $index :: Expected direction value to be $expectedValue but found $value" + ) + } + } + } + } + } + } + } + + fun parseAndAssertForAddFieldsStagesWithComputedValue( + query: PsiExpression, + expectedAddFieldsStagesCounts: Int, + addFieldsStagesExpectation: List>>> + ) { + SpringCriteriaDialectParser.parse(query).assert(IsCommand.CommandType.AGGREGATE) { + component { + assertEquals(HasSourceDialect.DialectName.SPRING_CRITERIA, name) + } + + collection> { + assertEquals("book", collection) + } + + component> { + assertEquals( + expectedAddFieldsStagesCounts, + children.size, + "Expected $expectedAddFieldsStagesCounts aggregation stages but found ${children.size}" + ) + } + + addFieldsStagesExpectation.forEachIndexed { + stageIndex, + (expectedAddedFieldsCount, addedFieldsExpectations) + -> + stageN(stageIndex, Name.ADD_FIELDS) { + component> { + assertEquals( + expectedAddedFieldsCount, + children.size, + "StageIndex $stageIndex :: Expected $expectedAddedFieldsCount added fields but found ${children.size}" + ) + } + + addedFieldsExpectations.forEachIndexed { + index, + (expectedFieldName, expectedValue) + -> + addedFieldN(index, Name.ADD_FIELDS, stageIndex) { + field> { + assertEquals( + expectedFieldName, + fieldName, + "StageIndex $stageIndex, AddedFieldIndex: $index :: Expected field with name $expectedFieldName but found field with name $fieldName" + ) + } + value> { + type.expression.field> { + assertEquals( + expectedValue, + fieldName, + "StageIndex $stageIndex, AddedFieldIndex: $index :: Expected computed value to have a field reference on $expectedValue but found $fieldName" + ) + + assertTrue( + displayName.startsWith("$"), + "StageIndex $stageIndex, AddedFieldIndex: $index :: Expected displayName of field in computed value reference to be \$$expectedValue but found $displayName" + ) + + assertEquals( + expectedValue, + fieldName, + "StageIndex $stageIndex, AddedFieldIndex: $index :: Expected computed value to have a field reference on $expectedValue but found $fieldName" + ) + } + } + } + } + } + } + } + } + } +}