2222import java .util .LinkedHashMap ;
2323import java .util .List ;
2424import java .util .Map ;
25- import java .util .function .BiFunction ;
26- import java .util .function .Consumer ;
2725
2826import javax .lang .model .element .Modifier ;
2927
3028import org .apache .commons .logging .Log ;
3129import org .apache .commons .logging .LogFactory ;
3230import org .jspecify .annotations .Nullable ;
31+
3332import org .springframework .aot .generate .ClassNameGenerator ;
3433import org .springframework .aot .generate .Generated ;
3534import org .springframework .data .projection .ProjectionFactory ;
@@ -58,8 +57,8 @@ class AotRepositoryBuilder {
5857 private final ProjectionFactory projectionFactory ;
5958 private final AotRepositoryFragmentMetadata generationMetadata ;
6059
61- private @ Nullable Consumer < AotRepositoryConstructorBuilder > constructorCustomizer ;
62- private @ Nullable BiFunction < Method , RepositoryInformation , @ Nullable MethodContributor <? extends QueryMethod >> methodContributorFunction ;
60+ private @ Nullable ConstructorCustomizer constructorCustomizer ;
61+ private @ Nullable MethodContributorFactory methodContributorFactory ;
6362 private ClassCustomizer customizer ;
6463
6564 private AotRepositoryBuilder (RepositoryInformation repositoryInformation , String moduleName ,
@@ -109,23 +108,21 @@ public AotRepositoryBuilder withClassCustomizer(ClassCustomizer classCustomizer)
109108 * @param constructorCustomizer must not be {@literal null}.
110109 * @return {@code this}.
111110 */
112- public AotRepositoryBuilder withConstructorCustomizer (
113- Consumer <AotRepositoryConstructorBuilder > constructorCustomizer ) {
111+ public AotRepositoryBuilder withConstructorCustomizer (ConstructorCustomizer constructorCustomizer ) {
114112
115113 this .constructorCustomizer = constructorCustomizer ;
116114 return this ;
117115 }
118116
119117 /**
120- * Configure a {@link MethodContributor}.
118+ * Configure a {@link MethodContributor} factory .
121119 *
122- * @param methodContributorFunction must not be {@literal null}.
120+ * @param methodContributorFactory must not be {@literal null}.
123121 * @return {@code this}.
124122 */
125- public AotRepositoryBuilder withQueryMethodContributor (
126- BiFunction <Method , RepositoryInformation , @ Nullable MethodContributor <? extends QueryMethod >> methodContributorFunction ) {
123+ public AotRepositoryBuilder withQueryMethodContributor (MethodContributorFactory methodContributorFactory ) {
127124
128- this .methodContributorFunction = methodContributorFunction ;
125+ this .methodContributorFactory = methodContributorFactory ;
129126 return this ;
130127 }
131128
@@ -170,7 +167,7 @@ private MethodSpec buildConstructor() {
170167 generationMetadata );
171168
172169 if (constructorCustomizer != null ) {
173- constructorCustomizer .accept (constructorBuilder );
170+ constructorCustomizer .customize (constructorBuilder );
174171 }
175172
176173 return constructorBuilder .buildConstructor ();
@@ -191,7 +188,8 @@ private AotRepositoryMetadata getAotRepositoryMetadata(List<AotRepositoryMethod>
191188 private void contributeMethod (Method method , RepositoryComposition repositoryComposition ,
192189 List <AotRepositoryMethod > methodMetadata , TypeSpec .Builder builder ) {
193190
194- if (repositoryInformation .isCustomMethod (method ) || (repositoryInformation .isBaseClassMethod (method ) && !repositoryInformation .isQueryMethod (method ))) {
191+ if (repositoryInformation .isCustomMethod (method )
192+ || (repositoryInformation .isBaseClassMethod (method ) && !repositoryInformation .isQueryMethod (method ))) {
195193
196194 RepositoryFragment <?> fragment = repositoryComposition .findFragment (method );
197195
@@ -205,9 +203,9 @@ private void contributeMethod(Method method, RepositoryComposition repositoryCom
205203 return ;
206204 }
207205
208- if (repositoryInformation .isQueryMethod (method ) && methodContributorFunction != null ) {
206+ if (repositoryInformation .isQueryMethod (method ) && methodContributorFactory != null ) {
209207
210- MethodContributor <? extends QueryMethod > contributor = methodContributorFunction . apply (method ,
208+ MethodContributor <? extends QueryMethod > contributor = methodContributorFactory . create (method ,
211209 repositoryInformation );
212210
213211 if (contributor != null ) {
@@ -273,16 +271,50 @@ public ProjectionFactory getProjectionFactory() {
273271 public interface ClassCustomizer {
274272
275273 /**
276- * Apply customization ot the AOT repository fragment class after it has been defined..
274+ * Apply customization ot the AOT repository fragment class after it has been defined.
277275 *
278- * @param information
279- * @param metadata
280- * @param builder
276+ * @param information repository information.
277+ * @param metadata metadata of the AOT repository fragment.
278+ * @param builder the actual builder.
281279 */
282280 void customize (RepositoryInformation information , AotRepositoryFragmentMetadata metadata , TypeSpec .Builder builder );
283281
284282 }
285283
284+ /**
285+ * Customizer interface to customize the AOT repository fragment constructor through
286+ * {@link AotRepositoryConstructorBuilder}.
287+ */
288+ public interface ConstructorCustomizer {
289+
290+ /**
291+ * Apply customization ot the AOT repository fragment constructor.
292+ *
293+ * @param constructorBuilder the builder to be customized.
294+ */
295+ void customize (AotRepositoryConstructorBuilder constructorBuilder );
296+
297+ }
298+
299+ /**
300+ * Factory interface to conditionally create {@link MethodContributor} instances. An implementation may decide whether
301+ * to return a {@link MethodContributor} or {@literal null}, if no method (code or metadata) should be contributed.
302+ */
303+ public interface MethodContributorFactory {
304+
305+ /**
306+ * Apply customization ot the AOT repository fragment constructor.
307+ *
308+ * @param method the method to be contributed.
309+ * @param information repository information.
310+ * @return the {@link MethodContributor} to be used. Can be {@literal null} if the method and method metadata should
311+ * not be contributed.
312+ */
313+ @ Nullable
314+ MethodContributor <? extends QueryMethod > create (Method method , RepositoryInformation information );
315+
316+ }
317+
286318 record AotBundle (JavaFile javaFile , AotRepositoryMetadata metadata ) {
287319 }
288320
0 commit comments