1818import java .lang .reflect .Parameter ;
1919import java .util .Arrays ;
2020import java .util .List ;
21+ import java .util .Optional ;
2122import java .util .regex .Pattern ;
2223import java .util .stream .Collectors ;
2324
2425import org .bson .Document ;
26+
27+ import org .springframework .core .annotation .MergedAnnotation ;
2528import org .springframework .data .mongodb .BindableMongoExpression ;
2629import org .springframework .data .mongodb .core .ExecutableFindOperation .FindWithQuery ;
2730import org .springframework .data .mongodb .core .ExecutableRemoveOperation .ExecutableRemove ;
3336import org .springframework .data .mongodb .repository .query .MongoQueryExecution .DeleteExecutionX .Type ;
3437import org .springframework .data .mongodb .repository .query .MongoQueryExecution .PagedExecution ;
3538import org .springframework .data .mongodb .repository .query .MongoQueryExecution .SlicedExecution ;
36- import org .springframework .data .repository .aot .generate .AotRepositoryMethodGenerationContext ;
39+ import org .springframework .data .mongodb .repository .query .MongoQueryMethod ;
40+ import org .springframework .data .repository .aot .generate .AotQueryMethodGenerationContext ;
3741import org .springframework .javapoet .ClassName ;
3842import org .springframework .javapoet .CodeBlock ;
3943import org .springframework .javapoet .CodeBlock .Builder ;
@@ -50,25 +54,29 @@ public class MongoBlocks {
5054
5155 private static final Pattern PARAMETER_BINDING_PATTERN = Pattern .compile ("\\ ?(\\ d+)" );
5256
53- static QueryBlockBuilder queryBlockBuilder (AotRepositoryMethodGenerationContext context ) {
54- return new QueryBlockBuilder (context );
57+ static QueryBlockBuilder queryBlockBuilder (AotQueryMethodGenerationContext context , MongoQueryMethod queryMethod ) {
58+ return new QueryBlockBuilder (context , queryMethod );
5559 }
5660
57- static QueryExecutionBlockBuilder queryExecutionBlockBuilder (AotRepositoryMethodGenerationContext context ) {
58- return new QueryExecutionBlockBuilder (context );
61+ static QueryExecutionBlockBuilder queryExecutionBlockBuilder (AotQueryMethodGenerationContext context ,
62+ MongoQueryMethod queryMethod ) {
63+ return new QueryExecutionBlockBuilder (context , queryMethod );
5964 }
6065
61- static DeleteExecutionBuilder deleteExecutionBlockBuilder (AotRepositoryMethodGenerationContext context ) {
62- return new DeleteExecutionBuilder (context );
66+ static DeleteExecutionBuilder deleteExecutionBlockBuilder (AotQueryMethodGenerationContext context ,
67+ MongoQueryMethod queryMethod ) {
68+ return new DeleteExecutionBuilder (context , queryMethod );
6369 }
6470
6571 static class DeleteExecutionBuilder {
6672
67- AotRepositoryMethodGenerationContext context ;
73+ private final AotQueryMethodGenerationContext context ;
74+ private final MongoQueryMethod queryMethod ;
6875 String queryVariableName ;
6976
70- public DeleteExecutionBuilder (AotRepositoryMethodGenerationContext context ) {
77+ public DeleteExecutionBuilder (AotQueryMethodGenerationContext context , MongoQueryMethod queryMethod ) {
7178 this .context = context ;
79+ this .queryMethod = queryMethod ;
7280 }
7381
7482 public DeleteExecutionBuilder referencing (String queryVariableName ) {
@@ -85,15 +93,16 @@ public CodeBlock build() {
8593 && !ObjectUtils .nullSafeEquals (TypeName .get (context .getRepositoryInformation ().getDomainType ()),
8694 context .getActualReturnType ());
8795
88- Object actualReturnType = isProjecting ? context .getActualReturnType ()
96+ Object actualReturnType = isProjecting ? context .getActualReturnType (). getType ()
8997 : context .getRepositoryInformation ().getDomainType ();
9098
9199 builder .add ("\n " );
92- builder .addStatement ("$T<$T> remover = $L.remove($T.class)" , ExecutableRemove .class , actualReturnType ,
100+ builder .addStatement ("$T<$T> remover = $L.remove($T.class)" , ExecutableRemove .class ,
101+ context .getRepositoryInformation ().getDomainType (),
93102 mongoOpsRef , context .getRepositoryInformation ().getDomainType ());
94103
95104 Type type = Type .FIND_AND_REMOVE_ALL ;
96- if (context . returnsSingleValue ()) {
105+ if (! queryMethod . isCollectionQuery ()) {
97106 if (!ClassUtils .isPrimitiveOrWrapper (context .getMethod ().getReturnType ())) {
98107 type = Type .FIND_AND_REMOVE_ONE ;
99108 } else {
@@ -103,7 +112,7 @@ public CodeBlock build() {
103112
104113 actualReturnType = ClassUtils .isPrimitiveOrWrapper (context .getMethod ().getReturnType ())
105114 ? ClassName .get (context .getMethod ().getReturnType ())
106- : context . returnsSingleValue () ? actualReturnType : context .getReturnType () ;
115+ : queryMethod . isCollectionQuery () ? context .getReturnTypeName () : actualReturnType ;
107116
108117 builder .addStatement ("return ($T) new $T(remover, $T.$L).execute($L)" , actualReturnType , DeleteExecutionX .class ,
109118 DeleteExecutionX .Type .class , type .name (), queryVariableName );
@@ -114,11 +123,14 @@ public CodeBlock build() {
114123
115124 static class QueryExecutionBlockBuilder {
116125
117- AotRepositoryMethodGenerationContext context ;
126+ private final AotQueryMethodGenerationContext context ;
127+ private final MongoQueryMethod queryMethod ;
118128 private String queryVariableName ;
129+ private boolean count , exists ;
119130
120- public QueryExecutionBlockBuilder (AotRepositoryMethodGenerationContext context ) {
131+ public QueryExecutionBlockBuilder (AotQueryMethodGenerationContext context , MongoQueryMethod queryMethod ) {
121132 this .context = context ;
133+ this .queryMethod = queryMethod ;
122134 }
123135
124136 QueryExecutionBlockBuilder referencing (String queryVariableName ) {
@@ -127,16 +139,24 @@ QueryExecutionBlockBuilder referencing(String queryVariableName) {
127139 return this ;
128140 }
129141
142+ QueryExecutionBlockBuilder count (boolean count ) {
143+ this .count = count ;
144+ return this ;
145+ }
146+
147+ QueryExecutionBlockBuilder exists (boolean exists ) {
148+ this .exists = exists ;
149+ return this ;
150+ }
151+
130152 CodeBlock build () {
131153
132154 String mongoOpsRef = context .fieldNameOf (MongoOperations .class );
133155
134156 Builder builder = CodeBlock .builder ();
135157
136- boolean isProjecting = context .getActualReturnType () != null
137- && !ObjectUtils .nullSafeEquals (TypeName .get (context .getRepositoryInformation ().getDomainType ()),
138- context .getActualReturnType ());
139- Object actualReturnType = isProjecting ? context .getActualReturnType ()
158+ boolean isProjecting = context .getReturnedType ().isProjecting ();
159+ Object actualReturnType = isProjecting ? context .getActualReturnType ().getType ()
140160 : context .getRepositoryInformation ().getDomainType ();
141161
142162 builder .add ("\n " );
@@ -150,24 +170,23 @@ CodeBlock build() {
150170 context .getRepositoryInformation ().getDomainType ());
151171 }
152172
153- String terminatingMethod = "all()" ;
154- if (context .returnsSingleValue ()) {
173+ String terminatingMethod ;
155174
156- if (context . returnsOptionalValue ()) {
157- terminatingMethod = "one ()" ;
158- } else if (context . isCountMethod () ) {
159- terminatingMethod = "count()" ;
160- } else if ( context . isExistsMethod ()) {
161- terminatingMethod = "exists()" ;
162- } else {
163- terminatingMethod = "oneValue()" ;
164- }
175+ if (queryMethod . isCollectionQuery () || queryMethod . isPageQuery () || queryMethod . isSliceQuery ()) {
176+ terminatingMethod = "all ()" ;
177+ } else if (count ) {
178+ terminatingMethod = "count()" ;
179+
180+ } else if ( exists ) {
181+ terminatingMethod = "exists()" ;
182+ } else {
183+ terminatingMethod = Optional . class . isAssignableFrom ( context . getReturnType (). toClass ()) ? "one()" : "oneValue()" ;
165184 }
166185
167- if (context . returnsPage ()) {
186+ if (queryMethod . isPageQuery ()) {
168187 builder .addStatement ("return new $T(finder, $L).execute($L)" , PagedExecution .class ,
169188 context .getPageableParameterName (), queryVariableName );
170- } else if (context . returnsSlice ()) {
189+ } else if (queryMethod . isSliceQuery ()) {
171190 builder .addStatement ("return new $T(finder, $L).execute($L)" , SlicedExecution .class ,
172191 context .getPageableParameterName (), queryVariableName );
173192 } else {
@@ -181,19 +200,22 @@ CodeBlock build() {
181200
182201 static class QueryBlockBuilder {
183202
184- AotRepositoryMethodGenerationContext context ;
203+ private final AotQueryMethodGenerationContext context ;
204+ private final MongoQueryMethod queryMethod ;
205+
185206 StringQuery source ;
186207 List <String > arguments ;
187208 private String queryVariableName ;
188209
189- public QueryBlockBuilder (AotRepositoryMethodGenerationContext context ) {
210+ public QueryBlockBuilder (AotQueryMethodGenerationContext context , MongoQueryMethod queryMethod ) {
190211 this .context = context ;
191212 this .arguments = Arrays .stream (context .getMethod ().getParameters ()).map (Parameter ::getName )
192213 .collect (Collectors .toList ());
193214
194215 // ParametersSource parametersSource = ParametersSource.of(repositoryInformation, metadata.getRepositoryMethod());
195216 // this.argumentSource = new MongoParameters(parametersSource, false);
196217
218+ this .queryMethod = queryMethod ;
197219 }
198220
199221 public QueryBlockBuilder filter (StringQuery query ) {
@@ -239,17 +261,20 @@ CodeBlock build() {
239261 }
240262
241263 String pageableParameter = context .getPageableParameterName ();
242- if (StringUtils .hasText (pageableParameter ) && !context . returnsPage () && !context . returnsSlice ()) {
264+ if (StringUtils .hasText (pageableParameter ) && !queryMethod . isPageQuery () && !queryMethod . isSliceQuery ()) {
243265 builder .addStatement ("$L.with($L)" , queryVariableName , pageableParameter );
244266 }
245267
246- String hint = context .annotationValue (Hint .class , "value" );
268+ MergedAnnotation <Hint > hintAnnotation = context .getAnnotation (Hint .class );
269+ String hint = hintAnnotation .isPresent () ? hintAnnotation .getString ("value" ) : null ;
247270
248271 if (StringUtils .hasText (hint )) {
249272 builder .addStatement ("$L.withHint($S)" , queryVariableName , hint );
250273 }
251274
252- String readPreference = context .annotationValue (ReadPreference .class , "value" );
275+ MergedAnnotation <ReadPreference > readPreferenceAnnotation = context .getAnnotation (ReadPreference .class );
276+ String readPreference = readPreferenceAnnotation .isPresent () ? readPreferenceAnnotation .getString ("value" ) : null ;
277+
253278 if (StringUtils .hasText (readPreference )) {
254279 builder .addStatement ("$L.withReadPreference($T.valueOf($S))" , queryVariableName ,
255280 com .mongodb .ReadPreference .class , readPreference );
0 commit comments