2222import java .util .LinkedHashMap ;
2323import java .util .List ;
2424import java .util .Map ;
25+ import java .util .function .Consumer ;
2526
2627import javax .lang .model .element .Modifier ;
2728
4344import org .springframework .javapoet .MethodSpec ;
4445import org .springframework .javapoet .TypeName ;
4546import org .springframework .javapoet .TypeSpec ;
47+ import org .springframework .util .Assert ;
4648
4749/**
4850 * Builder for AOT repository fragments.
@@ -59,9 +61,9 @@ class AotRepositoryBuilder {
5961 private final ProjectionFactory projectionFactory ;
6062 private final AotRepositoryFragmentMetadata generationMetadata ;
6163
62- private @ Nullable ConstructorCustomizer constructorCustomizer ;
64+ private @ Nullable Consumer < AotRepositoryConstructorBuilder > constructorCustomizer ;
6365 private @ Nullable MethodContributorFactory methodContributorFactory ;
64- private ClassCustomizer customizer ;
66+ private Consumer < AotRepositoryClassBuilder > classCustomizer ;
6567
6668 private AotRepositoryBuilder (RepositoryInformation repositoryInformation , String moduleName ,
6769 ProjectionFactory projectionFactory ) {
@@ -76,7 +78,7 @@ private AotRepositoryBuilder(RepositoryInformation repositoryInformation, String
7678 .initializer ("$T.getLog($T.class)" , TypeName .get (LogFactory .class ), this .generationMetadata .getTargetTypeName ())
7779 .build ());
7880
79- this .customizer = (info , builder ) -> {};
81+ this .classCustomizer = (builder ) -> {};
8082 }
8183
8284 /**
@@ -93,14 +95,14 @@ public static AotRepositoryBuilder forRepository(RepositoryInformation informati
9395 }
9496
9597 /**
96- * Configure a {@link ClassCustomizer } customizer.
98+ * Configure a {@link AotRepositoryConstructorBuilder } customizer.
9799 *
98100 * @param classCustomizer must not be {@literal null}.
99101 * @return {@code this}.
100102 */
101- public AotRepositoryBuilder withClassCustomizer (ClassCustomizer classCustomizer ) {
103+ public AotRepositoryBuilder withClassCustomizer (Consumer < AotRepositoryClassBuilder > classCustomizer ) {
102104
103- this .customizer = classCustomizer ;
105+ this .classCustomizer = classCustomizer ;
104106 return this ;
105107 }
106108
@@ -110,7 +112,8 @@ public AotRepositoryBuilder withClassCustomizer(ClassCustomizer classCustomizer)
110112 * @param constructorCustomizer must not be {@literal null}.
111113 * @return {@code this}.
112114 */
113- public AotRepositoryBuilder withConstructorCustomizer (ConstructorCustomizer constructorCustomizer ) {
115+ public AotRepositoryBuilder withConstructorCustomizer (
116+ Consumer <AotRepositoryConstructorBuilder > constructorCustomizer ) {
114117
115118 this .constructorCustomizer = constructorCustomizer ;
116119 return this ;
@@ -162,7 +165,12 @@ public AotBundle build() {
162165 generationMetadata .getFields ().values ().forEach (builder ::addField );
163166
164167 // finally customize the file itself
165- this .customizer .customize (repositoryInformation , builder );
168+ this .classCustomizer .accept (customizer -> {
169+
170+ Assert .notNull (customizer , "ClassCustomizer must not be null" );
171+ customizer .customize (builder );
172+ });
173+
166174 JavaFile javaFile = JavaFile .builder (packageName (), builder .build ()).build ();
167175 AotRepositoryMetadata metadata = getAotRepositoryMetadata (methodMetadata );
168176
@@ -171,11 +179,11 @@ public AotBundle build() {
171179
172180 private MethodSpec buildConstructor () {
173181
174- AotRepositoryConstructorBuilder constructorBuilder = new AotRepositoryConstructorBuilder ( repositoryInformation ,
182+ RepositoryConstructorBuilder constructorBuilder = new RepositoryConstructorBuilder (
175183 generationMetadata );
176184
177185 if (constructorCustomizer != null ) {
178- constructorCustomizer .customize (constructorBuilder );
186+ constructorCustomizer .accept (constructorBuilder );
179187 }
180188
181189 return constructorBuilder .buildConstructor ();
@@ -213,8 +221,7 @@ private void contributeMethod(Method method, RepositoryComposition repositoryCom
213221
214222 if (repositoryInformation .isQueryMethod (method ) && methodContributorFactory != null ) {
215223
216- MethodContributor <? extends QueryMethod > contributor = methodContributorFactory .create (method ,
217- repositoryInformation );
224+ MethodContributor <? extends QueryMethod > contributor = methodContributorFactory .create (method );
218225
219226 if (contributor != null ) {
220227
@@ -273,20 +280,6 @@ public ProjectionFactory getProjectionFactory() {
273280 return projectionFactory ;
274281 }
275282
276- /**
277- * Customizer interface to customize the AOT repository fragment class after it has been defined.
278- */
279- public interface ClassCustomizer {
280-
281- /**
282- * Apply customization ot the AOT repository fragment class after it has been defined.
283- *
284- * @param information the repository information that is used for the AOT fragment.
285- * @param builder the class builder to be customized.
286- */
287- void customize (RepositoryInformation information , TypeSpec .Builder builder );
288-
289- }
290283
291284 /**
292285 * Customizer interface to customize the AOT repository fragment constructor through
@@ -313,12 +306,11 @@ public interface MethodContributorFactory {
313306 * Apply customization ot the AOT repository fragment constructor.
314307 *
315308 * @param method the method to be contributed.
316- * @param information repository information.
317309 * @return the {@link MethodContributor} to be used. Can be {@literal null} if the method and method metadata should
318310 * not be contributed.
319311 */
320312 @ Nullable
321- MethodContributor <? extends QueryMethod > create (Method method , RepositoryInformation information );
313+ MethodContributor <? extends QueryMethod > create (Method method );
322314
323315 }
324316
0 commit comments